Set netcdf chunk size of iteration dimension to 1

[ajsilveira]

When trying to simulate more than 700 thermodynamic states, I am getting the following error=

self._reporter.write_mixing_statistics(self._n_accepted_matrix, self._n_proposed_matrix, self._iteration)
File "/home/silveira/miniconda3/envs/myenv/lib/python3.6/site-packages/openmmtools-0.19.0-py3.6-linux-x86_64.egg/openmmtools/multistate/multistatereporter.py", line 943, in write_mixing_statistics
n_states)
File "netCDF4/_netCDF4.pyx", line 2437, in netCDF4._netCDF4.Dataset.createVariable
File "netCDF4/_netCDF4.pyx", line 3439, in netCDF4._netCDF4.Variable.__init__
File "netCDF4/_netCDF4.pyx", line 1638, in netCDF4._netCDF4._ensure_nc_success
RuntimeError: NetCDF: Bad chunk sizes.

The PDB is here = https://gist.github.com/ajsilveira/423ffbc447339de440d14aa68236ed82
The serialized openmm system in here = https://gist.github.com/ajsilveira/e0e94aa959e29b86176f593361ee6070
I am using the following script =

import logging
import mdtraj as md
import yank
import itertools
from simtk import openmm, unit
from simtk.openmm import app
import openmmtools
from simtk.openmm import XmlSerializer
from openmmtools import states, mcmc
from openmmtools.states import GlobalParameterState
from openmmtools.multistate import SAMSSampler, MultiStateReporter
from openmmtools.alchemy import AbsoluteAlchemicalFactory, AlchemicalRegion, AlchemicalState

class MyComposableState(GlobalParameterState):
     lambda_sterics = GlobalParameterState.GlobalParameter('lambda_sterics', standard_value=1.0)
     lambda_electrostatics = GlobalParameterState.GlobalParameter('lambda_electrostatics', standard_value=1.0)
     lambda_restraints = GlobalParameterState.GlobalParameter('lambda_restraints', standard_value=12/160)
     K_parallel = GlobalParameterState.GlobalParameter('K_parallel',
                                                     standard_value=6000*unit.kilojoules_per_mole/unit.nanometer**2)
     K_orthogonal_max = GlobalParameterState.GlobalParameter('K_orthogonal_max',
                                                 standard_value=100*unit.kilojoules_per_mole/unit.nanometer**2)
     K_orthogonal_min = GlobalParameterState.GlobalParameter('K_orthogonal_min',
                                                 standard_value=10*unit.kilojoules_per_mole/unit.nanometer**2)

yank.utils.config_root_logger(verbose=True, log_file_path=None)
logger = logging.getLogger(__name__)
logging.root.setLevel(logging.DEBUG)
logging.basicConfig(level=logging.DEBUG)

from yank.experiment import ExperimentBuilder
platform = ExperimentBuilder._configure_platform('CUDA', 'mixed')

pdb = app.PDBFile('porin-ligand.pdb')
with open('openmm_system.xml', 'r') as infile:
    openmm_system = XmlSerializer.deserialize(infile.read())
topology = md.Topology.from_openmm(pdb.topology)
ligand_atoms = topology.select('(residue 422)')

factory = AbsoluteAlchemicalFactory(consistent_exceptions=False, disable_alchemical_dispersion_correction=True)
alchemical_region = AlchemicalRegion(alchemical_atoms=ligand_atoms)
alchemical_system = factory.create_alchemical_system(openmm_system, alchemical_region)

thermodynamic_state_ref = states.ThermodynamicState(system=alchemical_system, temperature=310*unit.kelvin, pressure=1.0*unit.atmospheres)

sampler_state = states.SamplerState(positions=pdb.positions, box_vectors=pdb.topology.getPeriodicBoxVectors())
lambda_sterics =[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]#, 1.0]#, 1.0]#, 0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1, 0.0]
lambda_electrostatics=[1.0, 0.9, 0.79, 0.69, 0.58, 0.46, 0.32, 0.18]#, 0.05]#, 0.0]#, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
lambda_restraints=[ i/160 for i in range(12, 112)]
K_parallel=[6000*unit.kilojoules_per_mole/unit.nanometer**2 for i in range(len(lambda_restraints))]
K_orthogonal_max=[100*unit.kilojoules_per_mole/unit.nanometer**2 for i in range(len(lambda_restraints))]
K_orthogonal_min=[10*unit.kilojoules_per_mole/unit.nanometer**2 for i in range(len(lambda_restraints))]

lambda_sterics_2d=lambda_sterics*len(lambda_restraints)
lambda_electrostatics_2d=lambda_electrostatics*len(lambda_restraints)
lambda_restraints_2d = list(itertools.chain.from_iterable(itertools.repeat(x, len(lambda_sterics)) for x in lambda_restraints))
K_parallel_2d = list(itertools.chain.from_iterable(itertools.repeat(x, len(lambda_sterics)) for x in K_parallel))
K_orthogonal_max_2d = list(itertools.chain.from_iterable(itertools.repeat(x, len(lambda_sterics)) for x in K_orthogonal_max))
K_orthogonal_min_2d = list(itertools.chain.from_iterable(itertools.repeat(x, len(lambda_sterics)) for x in K_orthogonal_min))

protocol = {'lambda_sterics': lambda_sterics_2d,
            'lambda_electrostatics': lambda_electrostatics_2d,
            'lambda_restraints': lambda_restraints_2d,
            'K_parallel': K_parallel_2d,
            'K_orthogonal_max': K_orthogonal_max_2d,
            'K_orthogonal_min': K_orthogonal_min_2d}

my_composable_state = MyComposableState.from_system(alchemical_system)
compound_states = states.create_thermodynamic_state_protocol(thermodynamic_state_ref, protocol=protocol,
                                                             composable_states=[my_composable_state])

move = mcmc.LangevinDynamicsMove(timestep=4*unit.femtosecond, collision_rate= 1.0/unit.picoseconds,
                                 n_steps=500, reassign_velocities=False)
simulation = SAMSSampler(mcmc_moves=move, number_of_iterations=1)

analysis_particle_indices = topology.select('(protein and mass > 3.0) or (residue 422 and mass > 3.0)')
reporter = MultiStateReporter('alchemical_test.nc', checkpoint_interval=2000, analysis_particle_indices=analysis_particle_indices)
simulation.create(thermodynamic_states=compound_states,
                  sampler_states=[sampler_state],
                  storage=reporter)
simulation.extend(n_iterations=1)

[andrrizzi]

This may be related to this: https://www.unidata.ucar.edu/mailing_lists/archives/netcdfgroup/2014/msg00260.html

If I remember correctly, the MultiStateReporter tries to assign a chunksize to the netcdf variables that make sense by looking at the size of the numpy array, but there's a limit to the maximum dimension of the chunksize that can be created, and probably the heuristic that we use to determine the chunksize doesn't consider this limit.

[ajsilveira]

@andrrizzi update: The problem was that I was asking more than 2Gb for the variables accepted and proposed with
chunksizes = (self._checkpoint_interval, n_states, n_states). I was setting checkpoint_interval=2000, which yields
(800x800x2000)x8bytes = 10.24 Gb. Setting checkpoint_interval=50 works fine. @jchodera mentioned that there are
speed issues probably associated with the chunksizes and that we could try the default chunking scheme.
However here it is stated that

The size and shape of chunks for each individual variable are
determined at creation time by the size of each variable element and by the shape of the variable,
specified by the ordered list of its dimensions and the lengths of each dimension,
with special rules for unlimited dimensions

which is what openmmtools already does. Would it be useful if I compare those setting by simulating this system?

[jchodera]

We specify the chunk sizes in several places in terms of self._checkpoint_interval:

• openmmtools/openmmtools/multistate/multistatereporter.py

  Lines 941 to 943 in 79d8b7e

  	chunksizes=(self._checkpoint_interval,
  	n_states,
  	n_states)

• openmmtools/openmmtools/multistate/multistatereporter.py

  Line 949 in 79d8b7e

  chunksizes=(self._checkpoint_interval,

Instead, let's replace those chunksizes=(self._checkpoint_interval, with chunksizes=(1,.

[andrrizzi]

Instead, let's replace those chunksizes=(self._checkpoint_interval, with chunksizes=(1,.

+1.

[andrrizzi]

For the records, the performance issue is described in choderalab/yank#1157.