hillclimber 0.1.0a1__py3-none-any.whl → 0.1.0a3__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of hillclimber might be problematic. Click here for more details.

hillclimber/interfaces.py CHANGED
@@ -69,13 +69,52 @@ class CollectiveVariable(Protocol):
69
69
  ...
70
70
 
71
71
 
72
- class MetadynamicsBiasCollectiveVariable(Protocol):
73
- """Protocol for metadata associated with a bias in PLUMED."""
72
+ class BiasProtocol(Protocol):
73
+ """Protocol for individual bias potentials that act on collective variables.
74
+
75
+ Individual bias potentials (restraints, walls) operate on collective variables
76
+ and generate their own PLUMED commands via to_plumed(). These biases are added
77
+ to MetaDynamicsModel via the `actions` parameter.
78
+
79
+ This is distinct from MetadynamicsBias, which is used for configuring the
80
+ collective METAD command itself (via the `bias_cvs` parameter).
81
+ """
82
+
83
+ cv: CollectiveVariable
84
+
85
+ def to_plumed(self, atoms: ase.Atoms) -> list[str]:
86
+ """Generate PLUMED input strings for this bias.
87
+
88
+ Parameters
89
+ ----------
90
+ atoms : ase.Atoms
91
+ The atomic structure to use for generating the PLUMED string.
92
+
93
+ Returns
94
+ -------
95
+ list[str]
96
+ List of PLUMED command strings.
97
+ """
98
+ ...
99
+
100
+
101
+ class MetadynamicsBias(Protocol):
102
+ """Protocol for metadynamics bias configuration objects.
103
+
104
+ This protocol defines the configuration structure for per-CV metadynamics
105
+ parameters (sigma, grid settings, etc.). These configuration objects are
106
+ used in the `bias_cvs` parameter of MetaDynamicsModel to build the collective
107
+ METAD command.
108
+
109
+ This is distinct from BiasProtocol, which is for individual bias potentials
110
+ (restraints, walls) that generate their own PLUMED commands and are added
111
+ via the `actions` parameter.
112
+ """
74
113
 
75
114
  cv: CollectiveVariable
76
115
  sigma: float | None = None
77
- grid_min: float | None = None
78
- grid_max: float | None = None
116
+ grid_min: float | str | None = None
117
+ grid_max: float | str | None = None
79
118
  grid_bin: int | None = None
80
119
 
81
120
 
@@ -84,7 +123,8 @@ __all__ = [
84
123
  "AtomSelector",
85
124
  "PlumedGenerator",
86
125
  "CollectiveVariable",
87
- "MetadynamicsBiasCollectiveVariable",
126
+ "BiasProtocol",
127
+ "MetadynamicsBias",
88
128
  ]
89
129
 
90
130
  def interfaces() -> dict[str, list[str]]:
@@ -1,4 +1,5 @@
1
1
  import dataclasses
2
+ import typing as t
2
3
  from pathlib import Path
3
4
 
4
5
  import ase.units
@@ -7,26 +8,28 @@ import zntrack
7
8
  from hillclimber.calc import NonOverwritingPlumed
8
9
  from hillclimber.interfaces import (
9
10
  CollectiveVariable,
10
- MetadynamicsBiasCollectiveVariable,
11
11
  NodeWithCalculator,
12
12
  PlumedGenerator,
13
13
  )
14
14
 
15
15
 
16
16
  @dataclasses.dataclass
17
- class MetaDBiasCV(MetadynamicsBiasCollectiveVariable):
18
- """Metadynamics bias on a collective variable.
17
+ class MetadBias:
18
+ """Metadynamics bias configuration for a collective variable.
19
19
 
20
20
  Parameters
21
21
  ----------
22
22
  cv : CollectiveVariable
23
23
  The collective variable to bias.
24
24
  sigma : float, optional
25
- The width of the Gaussian potential, by default None.
25
+ The width of the Gaussian potential in the same units as the CV
26
+ (e.g., Å for distances, radians for angles), by default None.
26
27
  grid_min : float | str, optional
27
- The minimum value of the grid, by default None.
28
+ The minimum value of the grid in CV units (or PLUMED expression like "-pi"),
29
+ by default None.
28
30
  grid_max : float | str, optional
29
- The maximum value of the grid, by default None.
31
+ The maximum value of the grid in CV units (or PLUMED expression like "pi"),
32
+ by default None.
30
33
  grid_bin : int, optional
31
34
  The number of bins in the grid, by default None.
32
35
 
@@ -48,22 +51,32 @@ class MetaDynamicsConfig:
48
51
 
49
52
  This contains only the global parameters that apply to all CVs.
50
53
 
54
+ Units
55
+ -----
56
+ hillclimber uses ASE units throughout. The UNITS line in the PLUMED input tells
57
+ PLUMED to interpret all values in ASE units:
58
+ - Distances: Ångström (Å)
59
+ - Energies: electronvolt (eV) - including HEIGHT, SIGMA for energy-based CVs, etc.
60
+ - Time: femtoseconds (fs)
61
+ - Temperature: Kelvin (K)
62
+
51
63
  Parameters
52
64
  ----------
53
65
  height : float, optional
54
- The height of the Gaussian potential in kJ/mol, by default 1.0.
66
+ The height of the Gaussian potential in eV, by default 1.0.
55
67
  pace : int, optional
56
- The frequency of Gaussian deposition, by default 500.
68
+ The frequency of Gaussian deposition in MD steps, by default 500.
57
69
  biasfactor : float, optional
58
70
  The bias factor for well-tempered metadynamics, by default None.
59
71
  temp : float, optional
60
72
  The temperature of the system in Kelvin, by default 300.0.
61
73
  file : str, optional
62
74
  The name of the hills file, by default "HILLS".
63
- adaptive : str, optional
64
- The adaptive scheme to use, by default "NONE".
75
+ adaptive : t.Literal["GEOM", "DIFF"] | None, optional
76
+ The adaptive scheme to use, by default None.
77
+ If None, no ADAPTIVE parameter is written to PLUMED.
65
78
  flush : int | None
66
- The frequency of flushing the output files.
79
+ The frequency of flushing the output files in MD steps.
67
80
  If None, uses the plumed default.
68
81
 
69
82
  Resources
@@ -77,7 +90,7 @@ class MetaDynamicsConfig:
77
90
  biasfactor: float | None = None
78
91
  temp: float = 300.0
79
92
  file: str = "HILLS"
80
- adaptive: str = "NONE" # NONE, DIFF, GEOM
93
+ adaptive: t.Literal["GEOM", "DIFF"] | None = None
81
94
  flush: int | None = None
82
95
 
83
96
 
@@ -115,7 +128,7 @@ class MetaDynamicsModel(zntrack.Node, NodeWithCalculator):
115
128
  ... x2=pn.SMARTSSelector(pattern="CO[C:1]"),
116
129
  ... prefix="d",
117
130
  ... )
118
- >>> metad_cv1 = pn.MetaDBiasCV(
131
+ >>> metad_cv1 = pn.MetadBias(
119
132
  ... cv=cv1, sigma=0.1, grid_min=0.0, grid_max=2.0, grid_bin=200
120
133
  ... )
121
134
  >>> model = pn.MetaDynamicsModel(
@@ -133,7 +146,7 @@ class MetaDynamicsModel(zntrack.Node, NodeWithCalculator):
133
146
  config: MetaDynamicsConfig = zntrack.deps()
134
147
  data: list[ase.Atoms] = zntrack.deps()
135
148
  data_idx: int = zntrack.params(-1)
136
- bias_cvs: list[MetaDBiasCV] = zntrack.deps(default_factory=list)
149
+ bias_cvs: list[MetadBias] = zntrack.deps(default_factory=list)
137
150
  actions: list[PlumedGenerator] = zntrack.deps(default_factory=list)
138
151
  timestep: float = zntrack.params(1.0) # in fs, default is 1 fs
139
152
  model: NodeWithCalculator = zntrack.deps()
@@ -188,8 +201,13 @@ class MetaDynamicsModel(zntrack.Node, NodeWithCalculator):
188
201
 
189
202
  sigmas, grid_mins, grid_maxs, grid_bins = [], [], [], []
190
203
 
204
+ # PLUMED UNITS line specifies conversion factors from ASE units to PLUMED's native units:
205
+ # - LENGTH=A: ASE uses Ångström (A), PLUMED native is nm → A is a valid PLUMED unit
206
+ # - TIME: ASE uses fs, PLUMED native is ps → 1 fs = 0.001 ps
207
+ # - ENERGY: ASE uses eV, PLUMED native is kJ/mol → 1 eV = 96.485 kJ/mol
208
+ # See: https://www.plumed.org/doc-master/user-doc/html/ (MD engine integration docs)
191
209
  plumed_lines.append(
192
- f"UNITS LENGTH=A TIME={1 / (1000 * ase.units.fs)} ENERGY={ase.units.mol / ase.units.kJ}"
210
+ f"UNITS LENGTH=A TIME={1/1000} ENERGY={ase.units.mol / ase.units.kJ}"
193
211
  )
194
212
 
195
213
  for bias_cv in self.bias_cvs:
@@ -197,17 +215,20 @@ class MetaDynamicsModel(zntrack.Node, NodeWithCalculator):
197
215
  plumed_lines.extend(cv_str)
198
216
  all_labels.extend(labels)
199
217
 
200
- # Collect per-CV parameters for later
201
- sigmas.append(str(bias_cv.sigma) if bias_cv.sigma is not None else None)
202
- grid_mins.append(
203
- str(bias_cv.grid_min) if bias_cv.grid_min is not None else None
204
- )
205
- grid_maxs.append(
206
- str(bias_cv.grid_max) if bias_cv.grid_max is not None else None
207
- )
208
- grid_bins.append(
209
- str(bias_cv.grid_bin) if bias_cv.grid_bin is not None else None
210
- )
218
+ # Collect per-CV parameters for later - repeat for each label
219
+ # PLUMED requires one parameter value per ARG, so if a CV generates
220
+ # multiple labels, we need to repeat the parameter values
221
+ for _ in labels:
222
+ sigmas.append(str(bias_cv.sigma) if bias_cv.sigma is not None else None)
223
+ grid_mins.append(
224
+ str(bias_cv.grid_min) if bias_cv.grid_min is not None else None
225
+ )
226
+ grid_maxs.append(
227
+ str(bias_cv.grid_max) if bias_cv.grid_max is not None else None
228
+ )
229
+ grid_bins.append(
230
+ str(bias_cv.grid_bin) if bias_cv.grid_bin is not None else None
231
+ )
211
232
 
212
233
  metad_parts = [
213
234
  "METAD",
@@ -216,16 +237,31 @@ class MetaDynamicsModel(zntrack.Node, NodeWithCalculator):
216
237
  f"PACE={self.config.pace}",
217
238
  f"TEMP={self.config.temp}",
218
239
  f"FILE={self.config.file}",
219
- f"ADAPTIVE={self.config.adaptive}",
220
240
  ]
241
+ if self.config.adaptive is not None:
242
+ metad_parts.append(f"ADAPTIVE={self.config.adaptive}")
221
243
  if self.config.biasfactor is not None:
222
244
  metad_parts.append(f"BIASFACTOR={self.config.biasfactor}")
223
245
 
224
246
  # Add SIGMA, GRID_MIN, GRID_MAX, GRID_BIN only if any value is set
225
247
  if any(v is not None for v in sigmas):
226
- metad_parts.append(
227
- f"SIGMA={','.join(v if v is not None else '0.0' for v in sigmas)}"
228
- )
248
+ # When using ADAPTIVE, PLUMED requires only one sigma value
249
+ if self.config.adaptive is not None:
250
+ # Validate that all sigma values are the same when adaptive is set
251
+ unique_sigmas = set(v for v in sigmas if v is not None)
252
+ if len(unique_sigmas) > 1:
253
+ raise ValueError(
254
+ f"When using ADAPTIVE={self.config.adaptive}, all CVs must have the same sigma value. "
255
+ f"Found different sigma values: {unique_sigmas}"
256
+ )
257
+ # Use the first non-None sigma value
258
+ sigma_value = next(v for v in sigmas if v is not None)
259
+ metad_parts.append(f"SIGMA={sigma_value}")
260
+ else:
261
+ # Standard mode: one sigma per CV
262
+ metad_parts.append(
263
+ f"SIGMA={','.join(v if v is not None else '0.0' for v in sigmas)}"
264
+ )
229
265
  if any(v is not None for v in grid_mins):
230
266
  metad_parts.append(
231
267
  f"GRID_MIN={','.join(v if v is not None else '0.0' for v in grid_mins)}"
@@ -240,10 +276,50 @@ class MetaDynamicsModel(zntrack.Node, NodeWithCalculator):
240
276
  )
241
277
 
242
278
  plumed_lines.append(f"metad: {' '.join(metad_parts)}")
243
- # Temporary until https://github.com/zincware/ZnTrack/issues/936
244
- from hillclimber.actions import PrintCVAction
245
- lines = PrintCVAction(cvs=[x.cv for x in self.bias_cvs], stride=100).to_plumed(atoms)
246
- plumed_lines.extend(lines)
279
+
280
+ # Track defined commands to detect duplicates and conflicts
281
+ # Map label -> full command for labeled commands (e.g., "d: DISTANCE ...")
282
+ defined_commands = {}
283
+ for line in plumed_lines:
284
+ # Check if this is a labeled command (format: "label: ACTION ...")
285
+ if ": " in line:
286
+ label = line.split(": ", 1)[0]
287
+ defined_commands[label] = line
288
+
289
+ # Add any additional actions (restraints, walls, print actions, etc.)
290
+ for action in self.actions:
291
+ action_lines = action.to_plumed(atoms)
292
+
293
+ # Filter out duplicate CV definitions, but detect conflicts
294
+ filtered_lines = []
295
+ for line in action_lines:
296
+ # Check if this is a labeled command
297
+ if ": " in line:
298
+ label = line.split(": ", 1)[0]
299
+
300
+ # Check if this label was already defined
301
+ if label in defined_commands:
302
+ # If the command is identical, skip (deduplication)
303
+ if defined_commands[label] == line:
304
+ continue
305
+ # If the command is different, raise error (conflict)
306
+ else:
307
+ raise ValueError(
308
+ f"Conflicting definitions for label '{label}':\n"
309
+ f" Already defined: {defined_commands[label]}\n"
310
+ f" New definition: {line}"
311
+ )
312
+ else:
313
+ # New labeled command, track it
314
+ defined_commands[label] = line
315
+ filtered_lines.append(line)
316
+ else:
317
+ # Unlabeled command, always add
318
+ filtered_lines.append(line)
319
+
320
+ plumed_lines.extend(filtered_lines)
321
+
322
+ # Add FLUSH if configured
247
323
  if self.config.flush is not None:
248
324
  plumed_lines.append(f"FLUSH STRIDE={self.config.flush}")
249
325
  return plumed_lines
hillclimber/opes.py ADDED
@@ -0,0 +1,357 @@
1
+ """OPES (On-the-fly Probability Enhanced Sampling) methods.
2
+
3
+ This module provides classes for OPES enhanced sampling, a modern alternative
4
+ to traditional metadynamics with improved convergence properties.
5
+ """
6
+
7
+ import dataclasses
8
+ from pathlib import Path
9
+
10
+ import ase.units
11
+ import zntrack
12
+
13
+ from hillclimber.calc import NonOverwritingPlumed
14
+ from hillclimber.interfaces import (
15
+ CollectiveVariable,
16
+ NodeWithCalculator,
17
+ PlumedGenerator,
18
+ )
19
+
20
+
21
+ @dataclasses.dataclass
22
+ class OPESBias:
23
+ """OPES bias configuration for a collective variable.
24
+
25
+ Parameters
26
+ ----------
27
+ cv : CollectiveVariable
28
+ The collective variable to bias.
29
+ sigma : float | str, optional
30
+ Initial kernel width in CV units (e.g., Å for distances, radians for angles).
31
+ Use "ADAPTIVE" for automatic adaptation (recommended).
32
+ If numeric, specifies the initial width.
33
+ Default: "ADAPTIVE".
34
+
35
+ Resources
36
+ ---------
37
+ - https://www.plumed.org/doc-master/user-doc/html/OPES_METAD/
38
+
39
+ Notes
40
+ -----
41
+ The ADAPTIVE sigma option automatically adjusts kernel widths based on
42
+ CV fluctuations, which is usually the best choice. The automatic width
43
+ is measured every ADAPTIVE_SIGMA_STRIDE steps (default: 10×PACE).
44
+ """
45
+
46
+ cv: CollectiveVariable
47
+ sigma: float | str = "ADAPTIVE"
48
+
49
+
50
+ @dataclasses.dataclass
51
+ class OPESConfig:
52
+ """Configuration for OPES_METAD and OPES_METAD_EXPLORE.
53
+
54
+ OPES (On-the-fly Probability Enhanced Sampling) is a modern enhanced
55
+ sampling method that samples well-tempered target distributions.
56
+
57
+ Units
58
+ -----
59
+ hillclimber uses ASE units throughout. The UNITS line in the PLUMED input tells
60
+ PLUMED to interpret all values in ASE units:
61
+ - Distances: Ångström (Å)
62
+ - Energies: electronvolt (eV) - including BARRIER, SIGMA_MIN, etc.
63
+ - Time: femtoseconds (fs)
64
+ - Temperature: Kelvin (K)
65
+
66
+ Parameters
67
+ ----------
68
+ barrier : float
69
+ Highest free energy barrier to overcome (eV). This is the key
70
+ parameter that determines sampling efficiency.
71
+ pace : int, optional
72
+ Frequency of kernel deposition in MD steps (default: 500).
73
+ temp : float, optional
74
+ Temperature in Kelvin (default: 300.0). If -1, retrieved from MD engine.
75
+ explore_mode : bool, optional
76
+ If True, uses OPES_METAD_EXPLORE which estimates target distribution
77
+ directly (better exploration, slower reweighting convergence).
78
+ If False, uses OPES_METAD which estimates unbiased distribution
79
+ (faster convergence, less exploration). Default: False.
80
+ biasfactor : float, optional
81
+ Well-tempered gamma factor. If not specified, uses default behavior.
82
+ Set to inf for custom target distributions.
83
+ compression_threshold : float, optional
84
+ Merge kernels if closer than this threshold in sigma units (default: 1.0).
85
+ file : str, optional
86
+ File to store deposited kernels (default: "KERNELS").
87
+ adaptive_sigma_stride : int, optional
88
+ MD steps between adaptive sigma measurements. If not set, uses 10×PACE.
89
+ sigma_min : float, optional
90
+ Minimum allowable sigma value for adaptive sigma in CV units.
91
+ state_wfile : str, optional
92
+ State file for writing exact restart information.
93
+ state_rfile : str, optional
94
+ State file for reading restart information.
95
+ state_wstride : int, optional
96
+ Frequency of STATE file writing in number of kernel depositions.
97
+ walkers_mpi : bool, optional
98
+ Enable multiple walker mode with MPI communication (default: False).
99
+ calc_work : bool, optional
100
+ Calculate and output accumulated work (default: False).
101
+ flush : int, optional
102
+ Frequency of flushing output files in MD steps.
103
+
104
+ Resources
105
+ ---------
106
+ - https://www.plumed.org/doc-master/user-doc/html/OPES_METAD/
107
+ - https://www.plumed.org/doc-master/user-doc/html/OPES_METAD_EXPLORE/
108
+ - Invernizzi & Parrinello, J. Phys. Chem. Lett. 2020
109
+
110
+ Notes
111
+ -----
112
+ **When to use OPES_METAD vs OPES_METAD_EXPLORE:**
113
+
114
+ - OPES_METAD: Use when you want quick convergence of reweighted free energy.
115
+ Estimates unbiased distribution P(s).
116
+
117
+ - OPES_METAD_EXPLORE: Use for systems with unknown barriers or when testing
118
+ new CVs. Allows more exploration but slower reweighting convergence.
119
+ Estimates target distribution p^WT(s) directly.
120
+
121
+ Both methods converge to the same bias given enough time, but approach it
122
+ differently. OPES is more sensitive to degenerate CVs than standard METAD.
123
+ """
124
+
125
+ barrier: float # kJ/mol
126
+ pace: int = 500
127
+ temp: float = 300.0
128
+ explore_mode: bool = False # False=OPES_METAD, True=OPES_METAD_EXPLORE
129
+ biasfactor: float | None = None
130
+ compression_threshold: float = 1.0
131
+ file: str = "KERNELS"
132
+ adaptive_sigma_stride: int | None = None
133
+ sigma_min: float | None = None
134
+ state_wfile: str | None = None
135
+ state_rfile: str | None = None
136
+ state_wstride: int | None = None
137
+ walkers_mpi: bool = False
138
+ calc_work: bool = False
139
+ flush: int | None = None
140
+
141
+
142
+ class OPESModel(zntrack.Node, NodeWithCalculator):
143
+ """OPES (On-the-fly Probability Enhanced Sampling) model.
144
+
145
+ Implements OPES_METAD and OPES_METAD_EXPLORE enhanced sampling methods.
146
+ OPES samples well-tempered target distributions and provides better
147
+ convergence properties than traditional metadynamics.
148
+
149
+ Parameters
150
+ ----------
151
+ config : OPESConfig
152
+ Configuration for the OPES simulation.
153
+ data : list[ase.Atoms]
154
+ Input data for simulation.
155
+ data_idx : int, optional
156
+ Index of data to use (default: -1).
157
+ bias_cvs : list[OPESBias], optional
158
+ Collective variables to bias (default: []).
159
+ actions : list[PlumedGenerator], optional
160
+ Additional actions like restraints, walls, print (default: []).
161
+ timestep : float, optional
162
+ Timestep in fs (default: 1.0).
163
+ model : NodeWithCalculator
164
+ Underlying force field model.
165
+
166
+ Examples
167
+ --------
168
+ >>> import hillclimber as hc
169
+ >>>
170
+ >>> # Define collective variables
171
+ >>> phi = hc.TorsionCV(...)
172
+ >>> psi = hc.TorsionCV(...)
173
+ >>>
174
+ >>> # OPES configuration (standard mode)
175
+ >>> config = hc.OPESConfig(
176
+ ... barrier=40.0, # kJ/mol
177
+ ... pace=500,
178
+ ... temp=300.0,
179
+ ... explore_mode=False # Use OPES_METAD
180
+ ... )
181
+ >>>
182
+ >>> # Bias configuration with adaptive sigma
183
+ >>> bias1 = hc.OPESBias(cv=phi, sigma="ADAPTIVE")
184
+ >>> bias2 = hc.OPESBias(cv=psi, sigma="ADAPTIVE")
185
+ >>>
186
+ >>> # Create OPES model
187
+ >>> opes = hc.OPESModel(
188
+ ... config=config,
189
+ ... bias_cvs=[bias1, bias2],
190
+ ... data=data.frames,
191
+ ... model=force_field,
192
+ ... timestep=0.5
193
+ ... )
194
+ >>>
195
+ >>> # For exploration mode, set explore_mode=True
196
+ >>> explore_config = hc.OPESConfig(
197
+ ... barrier=40.0,
198
+ ... pace=500,
199
+ ... explore_mode=True # Use OPES_METAD_EXPLORE
200
+ ... )
201
+
202
+ Resources
203
+ ---------
204
+ - https://www.plumed.org/doc-master/user-doc/html/OPES_METAD/
205
+ - https://www.plumed.org/doc-master/user-doc/html/OPES_METAD_EXPLORE/
206
+ - https://www.plumed.org/doc-master/user-doc/html/masterclass-22-03.html
207
+ - Invernizzi & Parrinello, J. Phys. Chem. Lett. 2020
208
+
209
+ Notes
210
+ -----
211
+ **Output Components:**
212
+ OPES provides several diagnostic outputs (accessible via PrintAction):
213
+ - opes.bias: Instantaneous bias potential value
214
+ - opes.rct: Convergence indicator (should flatten at convergence)
215
+ - opes.zed: Normalization estimate (should stabilize)
216
+ - opes.neff: Effective sample size
217
+ - opes.nker: Number of compressed kernels
218
+ - opes.work: Accumulated work (if calc_work=True)
219
+
220
+ **Advantages over Metadynamics:**
221
+ - Better convergence properties
222
+ - Automatic variance adaptation
223
+ - Lower systematic error
224
+ - More sensitive to degenerate CVs (helps identify CV problems)
225
+ """
226
+
227
+ config: OPESConfig = zntrack.deps()
228
+ data: list[ase.Atoms] = zntrack.deps()
229
+ data_idx: int = zntrack.params(-1)
230
+ bias_cvs: list[OPESBias] = zntrack.deps(default_factory=list)
231
+ actions: list[PlumedGenerator] = zntrack.deps(default_factory=list)
232
+ timestep: float = zntrack.params(1.0)
233
+ model: NodeWithCalculator = zntrack.deps()
234
+
235
+ figures: Path = zntrack.outs_path(zntrack.nwd / "figures", independent=True)
236
+
237
+ def run(self):
238
+ self.figures.mkdir(parents=True, exist_ok=True)
239
+ for cv in self.bias_cvs:
240
+ img = cv.cv.get_img(self.data[self.data_idx])
241
+ img.save(self.figures / f"{cv.cv.prefix}.png")
242
+
243
+ def get_calculator(
244
+ self, *, directory: str | Path | None = None, **kwargs
245
+ ) -> NonOverwritingPlumed:
246
+ if directory is None:
247
+ raise ValueError("Directory must be specified for PLUMED input files.")
248
+ directory = Path(directory)
249
+ directory.mkdir(parents=True, exist_ok=True)
250
+
251
+ lines = self.to_plumed(self.data[self.data_idx])
252
+ # replace FILE= with f"FILE={directory}/" inside config
253
+ lines = [line.replace("FILE=", f"FILE={directory}/") for line in lines]
254
+
255
+ # Write plumed input file
256
+ with (directory / "plumed.dat").open("w") as file:
257
+ for line in lines:
258
+ file.write(line + "\n")
259
+
260
+ kT = ase.units.kB * self.config.temp
261
+
262
+ return NonOverwritingPlumed(
263
+ calc=self.model.get_calculator(directory=directory),
264
+ atoms=self.data[self.data_idx],
265
+ input=lines,
266
+ timestep=float(self.timestep * ase.units.fs),
267
+ kT=float(kT),
268
+ log=(directory / "plumed.log").as_posix(),
269
+ )
270
+
271
+ def to_plumed(self, atoms: ase.Atoms) -> list[str]:
272
+ """Generate PLUMED input string for the OPES model."""
273
+ # check for duplicate CV prefixes
274
+ cv_labels = set()
275
+ for bias_cv in self.bias_cvs:
276
+ if bias_cv.cv.prefix in cv_labels:
277
+ raise ValueError(f"Duplicate CV prefix found: {bias_cv.cv.prefix}")
278
+ cv_labels.add(bias_cv.cv.prefix)
279
+
280
+ plumed_lines = []
281
+ all_labels = []
282
+
283
+ sigmas = []
284
+
285
+ # PLUMED UNITS line specifies conversion factors from ASE units to PLUMED's native units:
286
+ # - LENGTH=A: ASE uses Ångström (A), PLUMED native is nm → A is a valid PLUMED unit
287
+ # - TIME: ASE uses fs, PLUMED native is ps → 1 fs = 0.001 ps
288
+ # - ENERGY: ASE uses eV, PLUMED native is kJ/mol → 1 eV = 96.485 kJ/mol
289
+ # See: https://www.plumed.org/doc-master/user-doc/html/ (MD engine integration docs)
290
+ plumed_lines.append(
291
+ f"UNITS LENGTH=A TIME={1/1000} ENERGY={ase.units.mol / ase.units.kJ}"
292
+ )
293
+
294
+ for bias_cv in self.bias_cvs:
295
+ labels, cv_str = bias_cv.cv.to_plumed(atoms)
296
+ plumed_lines.extend(cv_str)
297
+ all_labels.extend(labels)
298
+
299
+ # Collect sigma values
300
+ if isinstance(bias_cv.sigma, str):
301
+ sigmas.append(bias_cv.sigma)
302
+ else:
303
+ sigmas.append(str(bias_cv.sigma))
304
+
305
+ # Determine which OPES method to use
306
+ method_name = "OPES_METAD_EXPLORE" if self.config.explore_mode else "OPES_METAD"
307
+
308
+ # Build OPES command
309
+ opes_parts = [
310
+ f"opes: {method_name}",
311
+ f"ARG={','.join(all_labels)}",
312
+ f"PACE={self.config.pace}",
313
+ f"BARRIER={self.config.barrier}",
314
+ f"TEMP={self.config.temp}",
315
+ ]
316
+
317
+ # Add SIGMA (required parameter)
318
+ # If all sigmas are the same, use single value; otherwise comma-separated
319
+ if len(set(sigmas)) == 1:
320
+ opes_parts.append(f"SIGMA={sigmas[0]}")
321
+ else:
322
+ opes_parts.append(f"SIGMA={','.join(sigmas)}")
323
+
324
+ # Add FILE and COMPRESSION_THRESHOLD
325
+ opes_parts.append(f"FILE={self.config.file}")
326
+ opes_parts.append(f"COMPRESSION_THRESHOLD={self.config.compression_threshold}")
327
+
328
+ # Optional parameters
329
+ if self.config.biasfactor is not None:
330
+ opes_parts.append(f"BIASFACTOR={self.config.biasfactor}")
331
+ if self.config.adaptive_sigma_stride is not None:
332
+ opes_parts.append(f"ADAPTIVE_SIGMA_STRIDE={self.config.adaptive_sigma_stride}")
333
+ if self.config.sigma_min is not None:
334
+ opes_parts.append(f"SIGMA_MIN={self.config.sigma_min}")
335
+ if self.config.state_wfile is not None:
336
+ opes_parts.append(f"STATE_WFILE={self.config.state_wfile}")
337
+ if self.config.state_rfile is not None:
338
+ opes_parts.append(f"STATE_RFILE={self.config.state_rfile}")
339
+ if self.config.state_wstride is not None:
340
+ opes_parts.append(f"STATE_WSTRIDE={self.config.state_wstride}")
341
+ if self.config.walkers_mpi:
342
+ opes_parts.append("WALKERS_MPI")
343
+ if self.config.calc_work:
344
+ opes_parts.append("CALC_WORK")
345
+
346
+ plumed_lines.append(" ".join(opes_parts))
347
+
348
+ # Add any additional actions (restraints, walls, print actions, etc.)
349
+ for action in self.actions:
350
+ action_lines = action.to_plumed(atoms)
351
+ plumed_lines.extend(action_lines)
352
+
353
+ # Add FLUSH if configured
354
+ if self.config.flush is not None:
355
+ plumed_lines.append(f"FLUSH STRIDE={self.config.flush}")
356
+
357
+ return plumed_lines