fennol.md.initial

  1from pathlib import Path
  2
  3import numpy as np
  4import jax.numpy as jnp
  5
  6from flax.core import freeze, unfreeze
  7
  8from ..utils.io import last_xyz_frame
  9
 10
 11from ..models import FENNIX
 12
 13from ..utils.periodic_table import PERIODIC_TABLE_REV_IDX, ATOMIC_MASSES
 14from .utils import us
 15from ..utils import detect_topology,parse_cell
 16
 17
 18def load_model(simulation_parameters):
 19    model_file = simulation_parameters.get("model_file")
 20    """@keyword[fennol_md] model_file
 21    Path to the machine learning model file (.fnx format). Required parameter.
 22    Type: str, Required
 23    """
 24    model_file = Path(str(model_file).strip())
 25    if not model_file.exists():
 26        raise FileNotFoundError(f"model file {model_file} not found")
 27    else:
 28        graph_config = simulation_parameters.get("graph_config", {})
 29        """@keyword[fennol_md] graph_config
 30        Advanced graph configuration for model initialization.
 31        Default: {}
 32        """
 33        model = FENNIX.load(model_file, graph_config=graph_config)  # \
 34        print(f"# model_file: {model_file}")
 35
 36    if "energy_terms" in simulation_parameters:
 37        energy_terms = simulation_parameters["energy_terms"]
 38        if isinstance(energy_terms, str):
 39            energy_terms = energy_terms.split()
 40        model.set_energy_terms(energy_terms)
 41        print("# energy terms:", model.energy_terms)
 42
 43    return model
 44
 45
 46def load_system_data(simulation_parameters, fprec):
 47    ## LOAD SYSTEM CONFORMATION FROM FILES
 48    system_name = str(simulation_parameters.get("system_name", "system")).strip()
 49    """@keyword[fennol_md] system_name
 50    Name prefix for output files. If not specified, uses the xyz filename stem.
 51    Default: "system"
 52    """
 53    indexed = simulation_parameters.get("xyz_input/indexed", False)
 54    """@keyword[fennol_md] xyz_input/indexed
 55    Whether first column contains atom indices (Tinker format).
 56    Default: False
 57    """
 58    has_comment_line = simulation_parameters.get("xyz_input/has_comment_line", True)
 59    """@keyword[fennol_md] xyz_input/has_comment_line
 60    Whether file contains comment lines.
 61    Default: True
 62    """
 63    xyzfile = Path(simulation_parameters.get("xyz_input/file", system_name + ".xyz"))
 64    """@keyword[fennol_md] xyz_input/file
 65    Path to xyz/arc coordinate file. Required parameter.
 66    Type: str, Required
 67    """
 68    if not xyzfile.exists():
 69        raise FileNotFoundError(f"xyz file {xyzfile} not found")
 70    system_name = str(simulation_parameters.get("system_name", xyzfile.stem)).strip()
 71    symbols, coordinates, _ = last_xyz_frame(
 72        xyzfile, indexed=indexed, has_comment_line=has_comment_line
 73    )
 74    coordinates = coordinates.astype(fprec)
 75    species = np.array([PERIODIC_TABLE_REV_IDX[s] for s in symbols], dtype=np.int32)
 76    nat = species.shape[0]
 77
 78    ## GET MASS
 79    mass_Da = np.array(ATOMIC_MASSES, dtype=fprec)[species]
 80    deuterate = simulation_parameters.get("deuterate", False)
 81    """@keyword[fennol_md] deuterate
 82    Replace hydrogen masses with deuterium masses.
 83    Default: False
 84    """
 85    if deuterate:
 86        print("# Replacing all hydrogens with deuteriums")
 87        mass_Da[species == 1] *= 2.0
 88
 89    totmass_Da = mass_Da.sum()
 90
 91    mass = mass_Da.copy()
 92    hmr = simulation_parameters.get("hmr", 0)
 93    """@keyword[fennol_md] hmr
 94    Hydrogen mass repartitioning factor. 0 = no repartitioning.
 95    Default: 0
 96    """
 97    if hmr > 0:
 98        print(f"# Adding {hmr} Da to H masses and repartitioning on others for total mass conservation.")
 99        Hmask = species == 1
100        added_mass = hmr * Hmask.sum()
101        mass[Hmask] += hmr
102        wmass = mass[~Hmask]
103        mass[~Hmask] -= added_mass * wmass/ wmass.sum()
104
105        assert np.isclose(mass.sum(), totmass_Da), "Mass conservation failed"
106
107    # convert to internal units
108    mass = mass / us.DA
109
110    ### GET TEMPERATURE
111    temperature = np.clip(simulation_parameters.get("temperature", 300.0), 1.0e-6, None)
112    """@keyword[fennol_md] temperature
113    Target temperature in Kelvin.
114    Default: 300.0
115    """
116    kT = us.K_B * temperature 
117
118    ### GET TOTAL CHARGE
119    total_charge = simulation_parameters.get("total_charge", None)
120    """@keyword[fennol_md] total_charge
121    Total system charge for charged systems.
122    Default: None (interpreted as 0)
123    """
124    if total_charge is None:
125        total_charge = 0
126    else:
127        total_charge = int(total_charge)
128        print("# total charge: ", total_charge,"e")
129
130    ### ENERGY UNIT
131    energy_unit_str = simulation_parameters.get("energy_unit", "kcal/mol")
132    """@keyword[fennol_md] energy_unit
133    Energy unit for output. Common options: 'kcal/mol', 'eV', 'Ha', 'kJ/mol'.
134    Default: "kcal/mol"
135    """
136    energy_unit = us.get_multiplier(energy_unit_str)
137
138    ## SYSTEM DATA
139    system_data = {
140        "name": system_name,
141        "nat": nat,
142        "symbols": symbols,
143        "species": species,
144        "mass": mass,
145        "mass_Da": mass_Da,
146        "totmass_Da": totmass_Da,
147        "temperature": temperature,
148        "kT": kT,
149        "total_charge": total_charge,
150        "energy_unit": energy_unit,
151        "energy_unit_str": energy_unit_str,
152    }
153    input_flags = simulation_parameters.get("model_flags", [])
154    """@keyword[fennol_md] model_flags
155    Additional flags to pass to the model.
156    Default: []
157    """
158    flags = {f:None for f in input_flags}
159
160    ### Set boundary conditions
161    cell = simulation_parameters.get("cell", None)
162    """@keyword[fennol_md] cell
163    Unit cell vectors. Required for PBC. It is a sequence of floats:
164    - 9 floats: components of cell vectors [ax, ay, az, bx, by, bz, cx, cy, cz]
165    - 6 floats: lengths and angles [a, b, c, alpha, beta, gamma]
166    - 3 floats: lengths of cell vectors [a, b, c] (orthorhombic)
167    - 1 float: length of cell vectors (cubic cell)
168    Lengths are in Angstroms, angles in degrees.
169    Default: None
170    """
171    if cell is not None:
172        cell = parse_cell(cell).astype(fprec)
173        # cell = np.array(cell, dtype=fprec).reshape(3, 3)
174        reciprocal_cell = np.linalg.inv(cell)
175        volume = np.abs(np.linalg.det(cell))
176        print("# cell matrix:")
177        for l in cell:
178            print("# ", l)
179        # print(cell)
180        dens = (totmass_Da/volume) * (us.MOL/us.CM**3)
181        print("# density: ", dens.item(), " g/cm^3")
182        minimum_image = simulation_parameters.get("minimum_image", True)
183        """@keyword[fennol_md] minimum_image
184        Use minimum image convention for neighbor lists in periodic systems.
185        Default: True
186        """
187        estimate_pressure = simulation_parameters.get("estimate_pressure", False)
188        """@keyword[fennol_md] estimate_pressure
189        Calculate and print pressure during simulation.
190        Default: False
191        """
192        print("# minimum_image: ", minimum_image)
193
194        crystal_input = simulation_parameters.get("xyz_input/crystal", False)
195        """@keyword[fennol_md] xyz_input/crystal
196        Use crystal coordinates.
197        Default: False
198        """
199        if crystal_input:
200            coordinates = coordinates @ cell
201
202        pbc_data = {
203            "cell": cell,
204            "reciprocal_cell": reciprocal_cell,
205            "volume": volume,
206            "minimum_image": minimum_image,
207            "estimate_pressure": estimate_pressure,
208        }
209        if minimum_image:
210            flags["minimum_image"] = None
211    else:
212        pbc_data = None
213    system_data["pbc"] = pbc_data
214    system_data["initial_coordinates"] = coordinates.copy()
215
216    ### TOPOLOGY
217    topology_key = simulation_parameters.get("topology", None)
218    """@keyword[fennol_md] topology
219    Topology specification for molecular systems. Use "detect" for automatic detection.
220    Default: None
221    """
222    if topology_key is not None:
223        topology_key = str(topology_key).strip()
224        if topology_key.lower() == "detect":
225            topology = detect_topology(species,coordinates,cell=cell)
226            np.savetxt(system_name +".topo", topology+1, fmt="%d")
227            print("# Detected topology saved to", system_name + ".topo")
228        else:
229            assert Path(topology_key).exists(), f"Topology file {topology_key} not found"
230            topology = np.loadtxt(topology_key, dtype=np.int32)-1
231            assert topology.shape[1] == 2, "Topology file must have two columns (source, target)"
232            print("# Topology loaded from", topology_key)
233    else:
234        topology = None
235    
236    system_data["topology"] = topology
237
238    ### PIMD
239    nbeads = simulation_parameters.get("nbeads", None)
240    """@keyword[fennol_md] nbeads
241    Number of beads for Path Integral MD.
242    Default: None
243    """
244    nreplicas = simulation_parameters.get("nreplicas", None)
245    """@keyword[fennol_md] nreplicas
246    Number of replicas for independent replica simulations.
247    Default: None
248    """
249    if nbeads is not None:
250        nbeads = int(nbeads)
251        print("# nbeads: ", nbeads)
252        system_data["nbeads"] = nbeads
253        coordinates = np.repeat(coordinates[None, :, :], nbeads, axis=0).reshape(-1, 3)
254        species = np.repeat(species[None, :], nbeads, axis=0).reshape(-1)
255        bead_index = np.arange(nbeads, dtype=np.int32).repeat(nat)
256        natoms = np.array([nat] * nbeads, dtype=np.int32)
257
258        eigmat = np.zeros((nbeads, nbeads))
259        for i in range(nbeads - 1):
260            eigmat[i, i] = 2.0
261            eigmat[i, i + 1] = -1.0
262            eigmat[i + 1, i] = -1.0
263        eigmat[-1, -1] = 2.0
264        eigmat[0, -1] = -1.0
265        eigmat[-1, 0] = -1.0
266        omk, eigmat = np.linalg.eigh(eigmat)
267        omk[0] = 0.0
268        omk = (nbeads * kT / us.HBAR) * omk**0.5
269        for i in range(nbeads):
270            if eigmat[i, 0] < 0:
271                eigmat[i] *= -1.0
272        eigmat = jnp.asarray(eigmat, dtype=fprec)
273        system_data["omk"] = omk
274        system_data["eigmat"] = eigmat
275        nreplicas = None
276    elif nreplicas is not None:
277        nreplicas = int(nreplicas)
278        print("# nreplicas: ", nreplicas)
279        system_data["nreplicas"] = nreplicas
280        system_data["mass"] = np.repeat(mass[None, :], nreplicas, axis=0).reshape(-1)
281        system_data["species"] = np.repeat(species[None, :], nreplicas, axis=0).reshape(
282            -1
283        )
284        coordinates = np.repeat(coordinates[None, :, :], nreplicas, axis=0).reshape(
285            -1, 3
286        )
287        species = np.repeat(species[None, :], nreplicas, axis=0).reshape(-1)
288        bead_index = np.arange(nreplicas, dtype=np.int32).repeat(nat)
289        natoms = np.array([nat] * nreplicas, dtype=np.int32)
290    else:
291        system_data["nreplicas"] = 1
292        bead_index = np.array([0] * nat, dtype=np.int32)
293        natoms = np.array([nat], dtype=np.int32)
294
295    conformation = {
296        "species": species,
297        "coordinates": coordinates,
298        "batch_index": bead_index,
299        "natoms": natoms,
300        "total_charge": total_charge,
301    }
302    if cell is not None:
303        cell = cell[None, :, :]
304        reciprocal_cell = reciprocal_cell[None, :, :]
305        if nbeads is not None:
306            cell = np.repeat(cell, nbeads, axis=0)
307            reciprocal_cell = np.repeat(reciprocal_cell, nbeads, axis=0)
308        elif nreplicas is not None:
309            cell = np.repeat(cell, nreplicas, axis=0)
310            reciprocal_cell = np.repeat(reciprocal_cell, nreplicas, axis=0)
311        conformation["cells"] = cell
312        conformation["reciprocal_cells"] = reciprocal_cell
313
314    additional_keys = simulation_parameters.get("additional_keys", {})
315    """@keyword[fennol_md] additional_keys
316    Additional custom keys for model input.
317    Default: {}
318    """
319    for key, value in additional_keys.items():
320        conformation[key] = value
321    
322    conformation["flags"] = flags
323
324    return system_data, conformation
325
326
327def initialize_preprocessing(simulation_parameters, model, conformation, system_data):
328    nblist_verbose = simulation_parameters.get("nblist_verbose", False)
329    """@keyword[fennol_md] nblist_verbose
330    Print detailed neighbor list information.
331    Default: False
332    """
333    nblist_skin = simulation_parameters.get("nblist_skin", -1.0)
334    """@keyword[fennol_md] nblist_skin
335    Neighbor list skin distance in Angstroms.
336    Default: -1.0 (automatic)
337    """
338
339    ### CONFIGURE PREPROCESSING
340    preproc_state = unfreeze(model.preproc_state)
341    layer_state = []
342    for st in preproc_state["layers_state"]:
343        stnew = unfreeze(st)
344        if nblist_skin > 0:
345            stnew["nblist_skin"] = nblist_skin
346        if "nblist_mult_size" in simulation_parameters:
347            stnew["nblist_mult_size"] = simulation_parameters["nblist_mult_size"]
348            """@keyword[fennol_md] nblist_mult_size
349            Multiplier for neighbor list size.
350            Default: None
351            """
352        if "nblist_add_neigh" in simulation_parameters:
353            stnew["add_neigh"] = simulation_parameters["nblist_add_neigh"]
354            """@keyword[fennol_md] nblist_add_neigh
355            Additional neighbors to include in lists.
356            Default: None
357            """
358        layer_state.append(freeze(stnew))
359    preproc_state["layers_state"] = layer_state
360    preproc_state = freeze(preproc_state)
361
362    ## initial preprocessing
363    preproc_state = preproc_state.copy({"check_input": True})
364    preproc_state, conformation = model.preprocessing(preproc_state, conformation)
365
366    preproc_state = preproc_state.copy({"check_input": False})
367
368    if nblist_verbose:
369        graphs_keys = list(model._graphs_properties.keys())
370        print("# graphs_keys: ", graphs_keys)
371        print("# nblist state:", preproc_state)
372
373    ### print model
374    if simulation_parameters.get("print_model", False):
375        """@keyword[fennol_md] print_model
376        Print detailed model information at startup.
377        Default: False
378        """
379        print(model.summarize(example_data=conformation))
380
381    return preproc_state, conformation