dyn: rework DynWrapper constructor

miniff/dyn.py

-from .util import randsphere
+from .util import randsphere, lookup
 from .kernel import Cell, ScalarFunctionWrapper, sf_parameters, compute_images
 from .potentials import harmonic_repulsion_potential_family
+from .units import element_masses
 
 from scipy.integrate import solve_ivp
 from scipy.optimize import minimize
@@ -86,22 +87,39 @@ class DynWrapper:
         ----------
         wrapper : ScalarFunctionWrapper
             The wrapper to evolve in time.
-        masses : Iterable
-            Particle masses.
+        masses : {None, 'atomic', Iterable}
+            Particle masses: None means unity masses,
+            'atomic' stands for atomic masses.
         v0 : np.ndarray
             Starting velocities.
         """
         if v0 is None:
-            v0 = np.zeros_like(wrapper.sample.coordinates)
+            v0 = sf_parameters(np.zeros_like(wrapper.sample.coordinates), np.zeros_like(wrapper.sample.vectors))
+        elif isinstance(v0, np.ndarray):
+            v0 = sf_parameters(v0, np.zeros_like(wrapper.sample.vectors))
+        elif isinstance(v0, tuple):
+            v0 = sf_parameters(*v0)
         else:
-            v0 = np.asanyarray(v0)
-        if masses is None:
-            masses = np.ones(wrapper.sample.size, dtype=float)
+            raise ValueError(f"Unknown starting velocities: {v0}")
+        if isinstance(masses, np.ndarray):
+            pass
         else:
-            masses = np.asanyarray(masses)
+            if masses is None:
+                atomic_masses = np.ones(wrapper.sample.size, dtype=float)
+            else:
+                if masses == "atomic":
+                    masses = element_masses()
+                atomic_masses = lookup(wrapper.sample.values, masses)
+            masses = []
+            if wrapper.include["cartesian"]:
+                masses.append(atomic_masses)
+            if wrapper.include["vectors"]:
+                masses.append([.75 * sum(atomic_masses) / np.pi ** 2] * 3)
+            masses = np.concatenate(masses)
         self.wrapper = wrapper
 
-        self.state = self.wrapper.sample.copy(proto=DynCell).copy(cartesian_velocity=v0, masses=masses, time=0)
+        self.state = self.wrapper.sample.copy(proto=DynCell).copy(cartesian_velocity=v0.cartesian, masses=masses,
+                                                                  time=0)
 
         self.snapshots = []
         self._history = []
@@ -233,7 +251,7 @@ def relax(cell, potentials, rtn_history=False, normalize=True, prefer_parallel=N
     rtn_history : bool
         If True, returns intermediate atomic configurations.
     normalize : bool
-        Normalizes the cell at each step if True.
+        If True, normalizes the cell at each step.
     prefer_parallel : bool
         A flag to prefer parallel potential computations.
     driver : Callable

miniff/ml_util.py

@@ -2201,7 +2201,7 @@ class SDWorkflow(Workflow):
         return self.cutoff
 
     @requires_fields("potentials", "cells")
-    def setup_dynamics(self, masses=None, **kwargs):
+    def setup_dynamics(self, masses='atomic', **kwargs):
         """
         Sets up dynamics wrappers.
 
@@ -2220,7 +2220,7 @@ class SDWorkflow(Workflow):
         if masses is None:
             masses = element_masses()
         self.wrappers = list(
-            DynWrapper(ScalarFunctionWrapper(i, self.potentials, **kwargs), masses=lookup(i.values, masses) if masses else None)
+            DynWrapper(ScalarFunctionWrapper(i, self.potentials, **kwargs), masses=masses)
             for i in self.cells
         )
         return self.wrappers

miniff/test_dyn.py

-from . import kernel, potentials, dyn, util
+from . import kernel, potentials, dyn, util, units
 
 import numpy as np
 from numpy import testing
@@ -10,11 +10,18 @@ class H2test(TestCase):
     def setUpClass(cls) -> None:
         cls.system = dyn.DynWrapper(kernel.ScalarFunctionWrapper(
             kernel.Cell(np.eye(3) * 10, [(.455, .5, .5), (.545, .5, .5)], ["H", "H"]),
-            [potentials.sw2_potential_family(gauge_a=7.049556227, gauge_b=0.6022245584, a=1.8, p=4, q=0, epsilon=0.5, sigma=1)],
+            [potentials.sw2_potential_family(gauge_a=7.049556227, gauge_b=0.6022245584, a=1.8, p=4, q=0, epsilon=0.5,
+                                             sigma=1)],
             pbc=False,
         ))
         cls.d_eq = 1.122462
 
+    def test_defaults(self):
+        testing.assert_equal(self.system.state.cartesian_velocity, np.zeros((2, 3)))
+        testing.assert_equal(self.system.state.masses, np.ones(2))
+        system = dyn.DynWrapper(self.system.wrapper, masses='atomic')
+        testing.assert_equal(system.state.masses, [units.element_masses()['h']] * 2)
+
     def test_jac(self):
         x = self.system.x
         j = self.system.wrapper.g(x)