Allow intermediate lambda values; ndarrays in the input parser might …

…break if the number of coefficient is different

hymd/compute_dihedral_forces.f90

@@ -2,25 +2,19 @@ subroutine cdf(force, r, dipoles, transfer_matrix, box, a, b, c, d, coeff, dtype
   use dipole_reconstruction
   implicit none
 
-  real(4), dimension(:,:), intent(in out) :: force
-  real(4), dimension(:,:), intent(in) :: r
-  real(4), dimension(:,:,:), intent(in out) :: dipoles
-  real(4), dimension(:,:,:,:), intent(in out) :: transfer_matrix
-  real(8), dimension(:), intent(in) :: box
-  integer, dimension(:), intent(in) :: a
-  integer, dimension(:), intent(in) :: b
-  integer, dimension(:), intent(in) :: c
-  integer, dimension(:), intent(in) :: d
-  real(4), dimension(:,:,:), intent(in) :: coeff 
-  integer, dimension(:), intent(in) :: dtype
-  integer, dimension(:), intent(in) :: bb_index
-  integer, intent(in) :: dipole_flag
+  real(4), intent(in out) :: force(:,:)
+  real(4), intent(in) :: r(:,:)
+  real(4), intent(in out) :: dipoles(:,:,:)
+  real(4), intent(in out) :: transfer_matrix(:,:,:,:)
+  real(8), intent(in) :: box(:)
+  integer, intent(in) :: a(:), b(:), c(:), d(:), dtype(:), bb_index(:), dipole_flag
+  real(4), intent(in) :: coeff(:,:,:)
   real(8), intent(out) :: energy
 
   integer :: ind, aa, bb, cc, dd, i
   integer, dimension(2) :: c_shape
   real(8), dimension(3) :: f, g, h, v, w, sc, fa, fb, fc, fd
-  real(8), dimension(5) :: c_v, c_k, d_v, d_k
+  real(8), dimension(5) :: c_v, c_k, c_coil, d_coil, d_v, d_k
   ! real(8), dimension(5) :: c_g, d_g
   real(8) :: energy_cbt, df_cbt
   real(8) :: force_const, eq_value
@@ -62,32 +56,28 @@ subroutine cdf(force, r, dipoles, transfer_matrix, box, a, b, c, d, coeff, dtype
     ! Cosine series, V_prop
     if (dtype(ind) == 0 .or. dtype(ind) == 1) then
       df_dih = 0.d0
-      ! Get shape of coeff to see how many arrays we have 
-      ! and use the shape to select the arrays
-      ! c_shape = shape(coeff(ind, :, :))
       c_v = coeff(ind, 1, :)
       d_v = coeff(ind, 2, :)
       call cosine_series(c_v, d_v, phi, energy, df_dih)
 
-      ! if (c_shape(1) > 4) then
-      !   print *, "Never for now"
-      !   ! c_shape = (6, whatever you provide)
-      !   c_v = coeff(ind, 3, :)
-      !   d_v = coeff(ind, 4, :)
-      !   call cosine_series(c_v, d_v, phi, energy, df_dih)
-      ! end if
+      c_coil = coeff(ind, 3, :)
+      d_coil = coeff(ind, 4, :)
+      if (count(c_coil == 0) /= size(c_coil) .and. &
+          count(d_coil == 0) /= size(d_coil)) then
+        call cosine_series(c_coil, d_coil, phi, energy, df_dih)
+      end if
     end if
 
     ! CBT potential
     if (dtype(ind) == 1) then
       ! V = V_prop + k * (gamma - gamma_0)**2      
 
-      c_k = coeff(ind, 3, :)
-      d_k = coeff(ind, 4, :)
+      c_k = coeff(ind, 5, :)
+      d_k = coeff(ind, 6, :)
 
       ! These are needed if gamma_0 is expressed as cosine series, not implemented
-      ! c_g = coeff(ind, 5, :)
-      ! d_g = phase(ind, 6, :)
+      ! c_g = coeff(ind, 7, :)
+      ! d_g = phase(ind, 8, :)
 
       call reconstruct( &
         f, r(bb, :), -g, box, c_k, d_k, phi, dipole_flag, &

hymd/force.py

@@ -40,7 +40,7 @@ class Dihedral:
     atom_2: str
     atom_3: str
     atom_4: str
-    coeffs: list
+    coeffs: np.ndarray
     # type: (0) Fourier or (1) CBT
     # Impropers to be specified in the toml?
     dih_type: int
@@ -216,10 +216,8 @@ def prepare_bonds(molecules, names, bonds, indices, config):
     bonds_4_atom3 = np.empty(len(bonds_4), dtype=int)
     bonds_4_atom4 = np.empty(len(bonds_4), dtype=int)
     # 4 => 2 sets of 2 parameters
-    # Might it be useful to decouple dihedral types to prevent having lots of zeros/empty slots?
-    # number_of_coeff = len(max([b[4] for b in bonds_4], key=len))
-    # len_of_coeff = len(max([cf for b in bonds_4 for cf in b[4]], key=len))
-    number_of_coeff = 4
+    # Might it be useful to decouple dihedral types to prevent having lots of zeros/empty arrays?
+    number_of_coeff = 6
     len_of_coeff = 5
     bonds_4_coeff = np.empty(
         (len(bonds_4), number_of_coeff, len_of_coeff), dtype=np.float64

hymd/input_parser.py

@@ -264,12 +264,13 @@ def propensity_potential_coeffs(x: float, comm):
         ]
     )
 
+    zero_add = np.zeros((2, 5))
     if x == -1:
-        return alpha_coeffs
+        return np.concatenate((alpha_coeffs, zero_add))
     elif x == 0:
-        return coil_coeffs
+        return np.concatenate((coil_coeffs, zero_add))
     elif x == 1:
-        return beta_coeffs
+        return np.concatenate((beta_coeffs, zero_add))
 
     abs_x = np.abs(x)
     if abs_x > 1:
@@ -361,22 +362,25 @@ def parse_config_toml(toml_content, file_path=None, comm=MPI.COMM_WORLD):
                     wrong_type_2 = len(b[1]) == 2 and not isinstance(b[1][0], list)
                     if wrong_len or wrong_type_1 or wrong_type_2:
                         err_str = (
-                            "The coefficients specified for the type 0 dihedral do not match the correct structure."
+                            "The coefficients specified for the dihedral type (0) do not match the correct structure."
                             + "Either use [lambda] or [[cn_prop], [dn_prop]], or select the correct dihedral type."
                         )
                         Logger.rank0.log(logging.ERROR, err_str)
                         if comm.Get_rank() == 0:
                             raise RuntimeError(err_str)
 
+                # FIXME: this is messy af, I don't like it
                 if dih_type == 0 and isinstance(b[1][0], (float, int)):
-                    coeff = propensity_potential_coeffs(b[1][0], comm).tolist()
+                    coeff = propensity_potential_coeffs(b[1][0], comm)
                 elif dih_type == 1 and len(b[1]) == 3:
-                    coeff = (
+                    coeff = np.array(
                         propensity_potential_coeffs(b[1][0][0], comm).tolist()
                         + b[1][1:]
                     )
+                elif dih_type == 2:
+                    coeff = np.array(b[1])
                 else:
-                    coeff = b[1]
+                    coeff = np.insert(np.array(b[1]), 2, np.zeros((2, 5)), axis=0)
 
                 config_dict["dihedrals"][i] = Dihedral(
                     atom_1=b[0][0],