Add Gaurav's simulations as an engine#91
Conversation
Add kill switch to simulation.
add full import path
Comment out kill switch check
SamTov
left a comment
There was a problem hiding this comment.
It looks good to me; I don't have a good understanding of the simulation side of things. Depending on how broad it will be, we could update the documentation side, but I would leave this for Gaurav.
I had a few comments regarding old code and the use of numba. I would prefer to just use Jax as it will also then compile nicely with the other parts of SwarmRL which also compile with Jax. I think it will also reduce the number of large dependencies. But if there is a real reason for it then I am open.
The other larger point is the use of a new Model. I think it would also be possible to reuse our ForceFunction parent class or, as I write in the comment, make a slightly more general parent class and start branching off with Espresso and MPI children.
| ) | ||
|
|
||
|
|
||
| @numba.jit |
There was a problem hiding this comment.
Is there a reason to with numba here and not Jax? If one requirement already includes GPU support and JIT compile I would lean towards sticking with that unless there is a reason not to.
| phi_i = state[i, 2] - r_ij_angle | ||
| phi_ij = phi_i | ||
|
|
||
| # f_mag_on_prefactor = ( |
There was a problem hiding this comment.
Do we need the old code here?
|
|
||
| from swarmrl.engine import Engine |
There was a problem hiding this comment.
I added the full import path.
| ) | ||
| for _ in range(n_slices): | ||
| # Check for a model termination condition. | ||
| if model.kill_switch: |
There was a problem hiding this comment.
I added a kill_switch check.
| offsets: np.array | ||
|
|
||
|
|
||
| class Model: |
There was a problem hiding this comment.
Is it possible to have these inherit from our ForceFunction? It sees like most of the code can be reused we should just rename colloids to particles or something along those lines. Perhaps we make a more general parent for the force functions so that we can have an ESPResSo and newer ones which then get their own custom type hints and we can do some checks when users try to use espresso actions in an MPI simulation. What do you think?
| round(self.params.time_slice / self.params.snapshot_interval) | ||
| ) | ||
| for _ in range(n_slices): | ||
| # Check for a model termination condition. |
There was a problem hiding this comment.
I added a kill switch here but commented it out as without the main parent classes for the force model it will break the tests.
Done:
To Do:
UnitsFuture: