Writers

Writers record training statistics at each step. JaQMC ships with console, CSV, and HDF5 writers. See Writers for background on output files and customizing console output.

Configuration

For writer config keys, see the configuration reference: Molecule, Solid, or Hall.

Base class

class jaqmc.writer.base.Writer

Base class for statistics writers.

Note

Do all I/O setup (file creation, opening handles) in the open context manager, not in __init__. In distributed runs, multiple processes may share the same filesystem, and side effects in __init__ can cause resource conflicts.

open(working_dir, stage_name, initial_step=0)

Context manager for resource handling.

This method manages resource lifecycle and side effects, such as initializing files or other I/O operations.

When restoring from a checkpoint, initial_step indicates where training will resume. Writers that persist to files should truncate any data beyond this point so that stale entries from a previous (interrupted) run are discarded.

Parameters:

• working_dir (UPath | Path) – The directory where artifacts should be stored.

• stage_name (str) – Name of the current training stage. File-backed writers may use it when resolving their output path template.

• initial_step (int, default: 0) – The step from which training will resume. Data written for steps >= initial_step should be discarded.

static resolve_path_template(working_dir, path_template, stage_name)

Resolve a writer path template against the working directory.

The template may contain {stage}, which is replaced with the current stage name before the path is resolved.

Parameters:

• working_dir (UPath | Path) – Base directory for relative paths.

• path_template (str) – Absolute path or path relative to working_dir.

• stage_name (str) – Name of the current stage.

Return type:

UPath

Returns:

The resolved output path.

Raises:

ValueError – If stage_name is empty or the template uses unsupported fields.

static to_scalar(val)

Returns a Python scalar from a JAX/NumPy scalar.

Return type:

Any

abstractmethod write(step, stats)

Write statistics for the current step.

Parameters:

• step (int) – Current iteration step.

• stats (Mapping[str, Any]) – Dictionary of statistics to write.

Return type:

None

class jaqmc.writer.base.Writers(writers=())

A collection of writers with master-process guarding.

Wraps multiple Writer instances and ensures that open and write only execute on the master process in distributed settings.

open(working_dir, stage_name, *, is_master=True, initial_step=0)

Open all writers on the master process.

Parameters:

• working_dir (UPath | Path) – Directory where artifacts should be stored.

• stage_name (str) – Name of the current training stage.

• is_master (bool, default: True) – Whether this is the master process.

• initial_step (int, default: 0) – The step from which training will resume. Data written for steps >= initial_step should be discarded.

Yields:

None.

write(step, stats)

Write statistics on the master process (no-op otherwise).

Parameters:

• step (int) – Current iteration step.

• stats (Mapping[str, Any]) – Dictionary of statistics to write.

Return type:

None

Built-in writers

class jaqmc.writer.console.ConsoleWriter(*, interval=1, fields='loss')

Bases: Writer

Writes statistics to the console via the standard logger.

Each field spec is [alias=]key[:format]:

• total_energy — display total_energy=...

• total_energy:.6f — with explicit format

• energy=total_energy — display as energy=...

• Lz=angular_momentum_z:+.4f — alias + format

Parameters:

• interval (int, default: 1) – Step interval for logging.

• fields (str, default: 'loss') – Comma-separated list of field specs.

open(working_dir, stage_name, initial_step=0)

Context manager for resource handling.

This method manages resource lifecycle and side effects, such as initializing files or other I/O operations.

When restoring from a checkpoint, initial_step indicates where training will resume. Writers that persist to files should truncate any data beyond this point so that stale entries from a previous (interrupted) run are discarded.

Parameters:

• working_dir – The directory where artifacts should be stored.

• stage_name – Name of the current training stage. File-backed writers may use it when resolving their output path template.

• initial_step (int, default: 0) – The step from which training will resume. Data written for steps >= initial_step should be discarded.

write(step, stats)

Write statistics for the current step.

Parameters:

• step (int) – Current iteration step.

• stats (Mapping[str, Any]) – Dictionary of statistics to write.

Return type:

None

class jaqmc.writer.console.FieldSpec(key, alias=None, fmt=None)

A single console output field.

Parameters:

• key (str) – Stat key to look up in the stats dictionary.

• alias (str | None, default: None) – Display name. Defaults to key.

• fmt (str | None, default: None) – Python format specifier (e.g. ".4f", "+.4f").

static parse(spec)

Parse [alias=]key[:format].

Stat keys may contain colons (e.g. energy:kinetic). The last :-separated segment is treated as a format specifier only when it contains a digit (e.g. .4f, +.4f).

Return type:

FieldSpec

Returns:

The parsed field specification.

Examples

>>> FieldSpec.parse("total_energy:.6f")
FieldSpec(key='total_energy', alias=None, fmt='.6f')

>>> FieldSpec.parse("Lz=angular_momentum_z:+.4f")
FieldSpec(key='angular_momentum_z', alias='Lz', fmt='+.4f')

Colons in stat keys are preserved (only the last segment is checked for format specifiers):

>>> FieldSpec.parse("energy:kinetic")
FieldSpec(key='energy:kinetic', alias=None, fmt=None)

class jaqmc.writer.csv.CSVWriter(*, path_template='{stage}_stats.csv')

Bases: Writer

Writes statistics to a CSV file.

Existing files are truncated to initial_step data rows upon open() before new rows are appended, so resumed runs discard stale rows past the restored checkpoint.

Parameters:

path_template (str, default: '{stage}_stats.csv') – Output path template. Relative paths are resolved under the working directory. The template may contain {stage}.

open(working_dir, stage_name, initial_step=0)

Context manager for resource handling.

This method manages resource lifecycle and side effects, such as initializing files or other I/O operations.

When restoring from a checkpoint, initial_step indicates where training will resume. Writers that persist to files should truncate any data beyond this point so that stale entries from a previous (interrupted) run are discarded.

Parameters:

• working_dir – The directory where artifacts should be stored.

• stage_name – Name of the current training stage. File-backed writers may use it when resolving their output path template.

• initial_step (int, default: 0) – The step from which training will resume. Data written for steps >= initial_step should be discarded.

write(step, stats)

Append one row of scalar statistics to the CSV file.

Parameters:

• step (int) – Current iteration step.

• stats (Mapping[str, Any]) – Statistics dictionary. Python scalars and scalar arrays are written; non-scalar values are skipped.

Raises:

ValueError – If called outside the open() context manager.

Return type:

None

class jaqmc.writer.hdf5.HDF5Writer(*, path_template='{stage}_stats.h5')

Bases: Writer

Writes statistics to an HDF5 file.

Existing files are truncated to initial_step data rows upon open() before new rows are appended, so resumed runs discard stale rows past the restored checkpoint.

Parameters:

path_template (str, default: '{stage}_stats.h5') – Output path template. Relative paths are resolved under the working directory. The template may contain {stage}.

open(working_dir, stage_name, initial_step=0)

Context manager for resource handling.

This method manages resource lifecycle and side effects, such as initializing files or other I/O operations.

When restoring from a checkpoint, initial_step indicates where training will resume. Writers that persist to files should truncate any data beyond this point so that stale entries from a previous (interrupted) run are discarded.

Parameters:

• working_dir – The directory where artifacts should be stored.

• stage_name – Name of the current training stage. File-backed writers may use it when resolving their output path template.

• initial_step (int, default: 0) – The step from which training will resume. Data written for steps >= initial_step should be discarded.

write(step, stats)

Write statistics for the current step.

Parameters:

• step (int) – Current iteration step.

• stats (Mapping[str, Any]) – Dictionary of statistics to write.

Return type:

None