Compute and save ensemble averages and Jacobians for each job

[mattwthompson]

Closes #145

• Compute ensemble averages in each job
• Compute Jacobians in each job
• Serialize ensemble averages in each job
• Serialize Jacobians in each job
• Add tests

[codecov-commenter]

Codecov Report

❌ Patch coverage is 29.49640% with 98 lines in your changes missing coverage. Please review.
✅ Project coverage is 80.03%. Comparing base (4db3654) to head (863c770).

Files with missing lines	Patch %	Lines
tyff/compute/_jacobian.py	0.00%	66 Missing ⚠️
tyff/compute/workflow.py	19.23%	21 Missing ⚠️
tyff/_serialization.py	76.31%	9 Missing ⚠️
tyff/compute/apps.py	60.00%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #173      +/-   ##
==========================================
- Coverage   81.40%   80.03%   -1.37%     
==========================================
  Files          56       58       +2     
  Lines        5070     5204     +134     
==========================================
+ Hits         4127     4165      +38     
- Misses        943     1039      +96     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Copilot]

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

[lilyminium]

Thanks Matt, this is a good start. I have several blocking comments; could you please also add a test?

[lilyminium]

I don't think this will work generally; the tensor_force_field is a series of tensors masquerading as a force field, and it builds a dense matrix of only the subset of parameters used in the interchanges passed to convert_interchange. Rebuilding the tensor_force_field for each property means the shapes of the jacobians returned will be different and un-alignable. I'd recommend passing in TensorSystem and TensorForceField as originally laid out in the function signature.

[lilyminium]

Could we please standardise the pressure unit across tyff? I think kPa is used elsewhere (e.g.

tyff/tyff/compute/_produce.py

Line 138 in 4db3654

pressure=pressure * openmm.unit.kilopascal,

) and I'm getting a bit confused. I'd probably vote for kPa

[mattwthompson]

Oddly this is not standardized across the framework, I will open a separate issue

$ grep --include="*.py" -r "\.atm" tyff | wc -l                                                                                                                        jacobian-app  ✱
      19
$ grep --include="*.py" -r "\.*pascal" tyff | wc -l                                                                                                                    jacobian-app  ✱
       4

[mattwthompson]

Running into an issue with the design here:

If we want the tyff.converters.convert_interchange(interchanges) call to be passed a list[Interchange] composing all of the unique molecules in the dataset, we will need to know all of the unique molecules in the dataset. But the current scope of workflow.submit_target¹ goes start to finish from system prep to the _get_ensemble_average_and_jacobian call. So when the dataset is more than trivially small, the first target is likely to hit the Jacobian step before the interchanges argument can properly be populated.

We could potentially separate the compute (from system prep to production MD runs, inclusive) from the Jacobian step such that the first Jacobian calculation is done after all production runs are done. This should work fine ... unless more targets are submitted.

So I'm thinking the Jacobians² should be computed in a step between the "compute" and before the prediction steps.

I'll have to come back to this later as it's clunky and I'd like something smoother.

Footnotes

1. If this method is the only one called, it's possible to assemble a list of unique molecules from the target_configs argument before firing off the compute apps. But if submit_target is used multiple times in one workflow, this is not feasible / a second call would invalidate the previously-saved Jacobians and require re-running the Jacobian app with potentially new molecules, each time a potentially new molecule is submitted ↩

2. And the steps prior: collecting the Interchanges for each molecule in the data set, the tensor force field, etc. ↩

[mattwthompson]

And now I'm getting hung up on how TensorSystems should be constructed inside of the Jacobian app. If this line takes in a list as long as the number of unique molecules in the training dataset:

    tensor_force_field, tensor_topologies = tyff.converters.convert_interchange(interchanges)

then should the n_copies argument to the TensorSystem constructor be sparse?

    system = tyff.TensorSystem(
        topologies=tensor_topologies,
        n_copies=n_copies,
        is_periodic=True,  # need to handle the case of gas simulations
    )

In other words, as I currently understand it, tensor_force_field and tensor_topologies are on the scale of the entire dataset (here 3 unique molecules, but in principle upwards of several hundred) but the TensorSystem is on the scale of a single job (almost always 1-2 unique molecules). Roughly L55-60 in the current tyff/compute/_jacobian.py has this mismatch

ipdb> n_molecules = compute_config["n_molecules"]
ipdb> n_copies = [int(n_molecules * x) for x in compute_config["x"]]
ipdb> n_copies
[200]
ipdb> len(tensor_topologies)
3

which is incongruous with the TensorSystem constructor expecting these to be of the same length (here 1 vs. 3 is a stand-in for the number of unique molecules in the target's chemical makeup vs. the entire dataset) in a long traceback ending at

File ~/software/tyff/tyff/potentials/_potentials.py:50, in broadcast_parameters(system, potential)
     48 parameters = []
     49 import ipdb; ipdb.set_trace()
---> 50 for topology, n_copies in zip(system.topologies, system.n_copies, strict=True):
     51     parameter_map = topology.parameters[potential.type]
     53     topology_parameters = parameter_map.assignment_matrix @ potential.parameters

ValueError: zip() argument 2 is shorter than argument 1

Put a slightly different way (if repetitive) here's a collection of the unique SMILES in the (canned) dataset (n=3) vs. the compute config for a single job (n=1)

ipdb> [interchange.topology.molecule(0).to_smiles(explicit_hydrogens=False) for interchange in interchanges]
['CCCOC=C', 'COCCO', 'CCC(=O)OC']
ipdb> compute_config
{'tag': 'liquid', 'force_field': 'openff-2.3.0.offxml', 'n_molecules': 200, 'replicate_index': 0, 'smiles': ['COCCO'], 'x': [1.0], 'temperature': 293.15, 'pressure': 101.3, 'density': 0.9648800000000002}

Maybe I just need to massage the data to make it, using the above values, look like n_copies=[0, 200, 0]? For larger datasets this might be an awkward-to-me [0, 0, 0, 0, 0, 0, 0, ..., 200, 0, 0, 0, 0, 0, 0, ..., 0, 0, 0, 0, 0]

[lilyminium]

n_copies

You can construct a TensorSystem from just the tensor topologies in the simulation, this is usually dense not sparse.

design

Let's discuss this synchronously tomorrow, but I thought about and tossed up some scenarios where:

a) we just pass in the TensorForceField / tensor topologies to workflow submit_target to pass into this method, as originally laid out in the issue code (but this probably makes tyff difficult to use for benchmarking)
b) we construct TensorForceField expansively to include all parameters of the FF in order, to fix alignment and the issue of new molecules (I think unique per-molecule charges makes this impossible)

At the end of the day I think we have to simply require that the user specifies the entire dataset up front. This is already kind of essential since we need a TensorForceField created up front to be fit (see the MVP issue). The driver could serialize that with the dump_tensor_force_field and pass it through submit_target (perhaps either in replacement of the current FF argument, or with some checks to make sure they have the same parameters?)

Inside the function what we could do instead is still create a local TensorForceField but have it remap from the global reference, and still take the Jacobian wrt the packed reference force field so autograd stil gets the right reference column.

def _gather_from_reference(
    local: tyff.TensorForceField, reference: tyff.TensorForceField
) -> tyff.TensorForceField:
    """Rebuild ``local`` with values indexed out of ``reference``, keeping local row order so the
    topologies' parameter maps stay valid."""
    reference_by_type = reference.potentials_by_type
    potentials = []

    for potential in local.potentials:
        ref = reference_by_type[potential.type]  # KeyError: potential type missing from reference
        assert potential.parameter_cols == ref.parameter_cols, potential.type

        ref_idx = {key: i for i, key in enumerate(ref.parameter_keys)}
        idx = torch.tensor([ref_idx[key] for key in potential.parameter_keys])  # KeyError: unseen parameter

        attributes = potential.attributes
        if potential.attribute_cols is not None:
            attr_idx = torch.tensor([ref.attribute_cols.index(col) for col in potential.attribute_cols])
            attributes = ref.attributes[attr_idx]

        potentials.append(dataclasses.replace(potential, parameters=ref.parameters[idx], attributes=attributes))

    v_sites = local.v_sites
    if v_sites is not None:
        ref_idx = {key: i for i, key in enumerate(reference.v_sites.keys)}
        idx = torch.tensor([ref_idx[key] for key in v_sites.keys])
        v_sites = dataclasses.replace(v_sites, parameters=reference.v_sites.parameters[idx])

    return tyff.TensorForceField(potentials, v_sites)

In the existing function, then:

local_ff, tensor_topologies = tyff.converters.convert_interchange(interchanges) # of just this system

reference_ff = ... # deserialize global reference TensorForceField

tensors, parameter_lookup, attribute_lookup, has_v_sites = _pack_force_field(reference_ff)  # was local
parameters = torch.cat([...]).requires_grad_(True)  # unchanged
...
reference_worker_ff = _unpack_force_field(tensors, parameter_lookup, attribute_lookup, has_v_sites, reference_ff)
worker_ff = _gather_from_reference(local_ff, reference_worker_ff)
means, _ = tyff.mm.compute_ensemble_averages(system, worker_ff, ...)

Lastly I also considered that we don't really need the Jacobian for benchmarking and that part can stay using a locally generated FF -- this is probably overengineered from that case, but also maybe this is not the right PR to address that.