Fix minimization and potential initialization

chiron/multistate.py

@@ -259,8 +259,18 @@ def _allocate_variables(self, thermodynamic_states, sampler_states):
         self._neighborhoods = np.zeros([self.n_replicas, self.n_states], "i1")
 
     def _minimize_replica(
-        self, replica_id: int, tolerance: unit.Quantity, max_iterations: int
+        self, replica_id: int, tolerance: unit.Quantity, max_iterations: int = 1_000
     ):
+        """
+        Minimizes the energy of a replica using the provided parameters.
+
+        Parameters
+        ----------
+        replica_id (int): The ID of the replica.
+        tolerance (unit.Quantity): The tolerance for convergence.
+        max_iterations (int, optional): The maximum number of iterations for the minimization. Defaults to 1_000.
+        """
+
         from chiron.minimze import minimize_energy
 
         # Retrieve thermodynamic and sampler states.
@@ -278,9 +288,9 @@ def _minimize_replica(
             sampler_state.x0,
             thermodynamic_state.potential.compute_energy,
             self.nbr_list,
-            maxiter=0,
+            maxiter=max_iterations,
         )
-        sampler_state.positions = results.params
+        sampler_state.x0 = results.params
         final_energy = thermodynamic_state.get_reduced_potential(sampler_state)
         log.debug(
             f"Replica {replica_id + 1}/{self.n_replicas}: final energy {final_energy:8.3f}kT"
@@ -289,7 +299,7 @@ def _minimize_replica(
     def minimize(
         self,
         tolerance=1.0 * unit.kilojoules_per_mole / unit.nanometers,
-        max_iterations=0,
+        max_iterations: int = 1_000,
     ):
         """Minimize all replicas.
 
@@ -314,16 +324,12 @@ def minimize(
         log.debug("Minimizing all replicas...")
 
         # minimization and update sampler states
-        minimized_positions, sampler_state_ids = [], []
         for replica_id in range(self.n_replicas):
             self._minimize_replica(replica_id, tolerance, max_iterations)
 
     def equilibrate(self, n_iterations, mcmc_moves=None):
         """Equilibrate all replicas.
 
-        This does not increase the iteration counter. The equilibrated
-        positions are stored at the end.
-
         Parameters
         ----------
         n_iterations : int

chiron/potential.py

@@ -272,9 +272,24 @@ def __init__(
         self,
         topology: Topology,
         k: unit.Quantity = 1.0 * unit.kilocalories_per_mole / unit.angstrom**2,
-        x0: unit.Quantity = 0.0 * unit.angstrom,
+        x0: unit.Quantity = [[0.0, 0.0, 0.0]] * unit.angstrom,
         U0: unit.Quantity = 0.0 * unit.kilocalories_per_mole,
     ):
+        """
+        Initialize a HarmonicOscillatorPotential object.
+
+        Parameters:
+        ----------
+        topology : Topology
+            The topology object representing the molecular system.
+        k : unit.Quantity, optional
+            The spring constant of the harmonic potential. Default is 1.0 kcal/mol/Å^2.
+        x0 : unit.Quantity, optional
+            The equilibrium position of the harmonic potential. Default is [0.0,0.0,0.0] Å.
+        U0 : unit.Quantity, optional
+            The offset potential energy of the harmonic potential. Default is 0.0 kcal/mol.
+        """
+
         if not isinstance(topology, Topology):
             if not isinstance(
                 topology, property
@@ -298,6 +313,7 @@ def __init__(
             raise ValueError(
                 f"x0 must be a unit.Quantity with units of distance, x0.unit = {x0.unit}"
             )
+        assert x0.shape[1] == 3, f"x0 must be a NX3 vector, x0.shape = {x0.shape}"
         if not U0.unit.is_compatible(unit.kilocalories_per_mole):
             raise ValueError(
                 f"U0 must be a unit.Quantity with units of energy, U0.unit = {U0.unit}"