diff --git a/.github/workflows/gromacs-integration.yaml b/.github/workflows/gromacs-integration.yaml index 4b85bfd..98209d4 100644 --- a/.github/workflows/gromacs-integration.yaml +++ b/.github/workflows/gromacs-integration.yaml @@ -50,6 +50,11 @@ jobs: install-tests: true installer: conda shell: bash -l {0} + + - name: Install package + run: | + python --version + python -m pip install . --no-deps - name: Python information run: | diff --git a/.github/workflows/lammps-integration.yaml b/.github/workflows/lammps-integration.yaml index 252cd7a..979e75a 100644 --- a/.github/workflows/lammps-integration.yaml +++ b/.github/workflows/lammps-integration.yaml @@ -41,7 +41,6 @@ jobs: show-channel-urls: true miniconda-version: latest - - name: Install MDAnalysis version uses: MDAnalysis/install-mdanalysis@main with: @@ -49,6 +48,11 @@ jobs: install-tests: true installer: conda shell: bash -l {0} + + - name: Install package + run: | + python --version + python -m pip install . --no-deps - name: Python information run: | diff --git a/.github/workflows/namd-integration.yaml b/.github/workflows/namd-integration.yaml index 7652630..d5dd157 100644 --- a/.github/workflows/namd-integration.yaml +++ b/.github/workflows/namd-integration.yaml @@ -51,6 +51,11 @@ jobs: installer: conda shell: bash -l {0} + - name: Install package + run: | + python --version + python -m pip install . --no-deps + - name: Python information run: | which python @@ -68,7 +73,7 @@ jobs: - name: Build namd run: | - cd namd/namd-3.0 + cd namd tar xf charm-8.0.0.tar cd charm-8.0.0 ./build charm++ ucx-linux-x86_64 ompipmix --with-production diff --git a/imdclient/IMDClient.py b/imdclient/IMDClient.py index 22c1e51..0b0e209 100644 --- a/imdclient/IMDClient.py +++ b/imdclient/IMDClient.py @@ -365,7 +365,6 @@ def run(self): logger.debug("IMDProducer: Stopping run loop") # Tell consumer not to expect more frames to be added self._buf.notify_producer_finished() - return def _expect_header(self, expected_type, expected_value=None): @@ -564,6 +563,12 @@ def _parse_imdframe(self): logger.debug( f"IMDProducer: Time: {self._imdf.time}, dt: {self._imdf.dt}, step: {self._imdf.step}" ) + if self._imdsinfo.energies: + self._expect_header(IMDHeaderType.IMD_ENERGIES, expected_value=1) + self._read(self._energies) + self._imdf.energies.update( + parse_energy_bytes(self._energies, self._imdsinfo.endianness) + ) if self._imdsinfo.box: self._expect_header(IMDHeaderType.IMD_BOX, expected_value=1) self._read(self._box) @@ -604,13 +609,6 @@ def _parse_imdframe(self): ).reshape((self._n_atoms, 3)), ) - if self._imdsinfo.energies: - self._expect_header(IMDHeaderType.IMD_ENERGIES, expected_value=1) - self._read(self._energies) - self._imdf.energies.update( - parse_energy_bytes(self._energies, self._imdsinfo.endianness) - ) - def __del__(self): logger.debug("IMDProducer: I am being deleted") diff --git a/imdclient/IMDREADER.py b/imdclient/IMDREADER.py index 7ae8215..21ed403 100644 --- a/imdclient/IMDREADER.py +++ b/imdclient/IMDREADER.py @@ -71,6 +71,7 @@ class IMDReader(ReaderBase): """ format = "IMD" + one_pass = True @store_init_arguments def __init__( diff --git a/imdclient/data/lammps_trj.h5md b/imdclient/data/lammps_trj.h5md index bfd16d9..51f5c28 100644 Binary files a/imdclient/data/lammps_trj.h5md and b/imdclient/data/lammps_trj.h5md differ diff --git a/imdclient/data/lammps_v3.in b/imdclient/data/lammps_v3.in index cba186a..8dcff3f 100644 --- a/imdclient/data/lammps_v3.in +++ b/imdclient/data/lammps_v3.in @@ -53,11 +53,11 @@ fix 1 all nve # Create source of truth trajectory # dump h5md1 all h5md 1 lammps_trj.h5md position velocity force box yes -# dump_modify h5md1 unwrap on +# dump_modify h5md1 unwrap no ## IMD settings # https://docs.lammps.org/fix_imd.html -fix 2 all imd 8888 version 3 unwrap on nowait off +fix 2 all imd 8888 version 3 unwrap off nowait off ## Run MD sim run 100 diff --git a/imdclient/tests/base.py b/imdclient/tests/base.py index ddf2e91..83845bb 100644 --- a/imdclient/tests/base.py +++ b/imdclient/tests/base.py @@ -65,7 +65,9 @@ def assert_allclose_with_logging(a, b, rtol=1e-07, atol=0, equal_nan=False): class IMDv3IntegrationTest: @pytest.fixture() - def run_sim_and_wait(self, command, match_string, simulation): + def run_sim_and_wait(self, tmp_path, command, match_string): + old_cwd = Path.cwd() + os.chdir(tmp_path) p = subprocess.Popen( command, stdout=subprocess.PIPE, @@ -87,14 +89,8 @@ def run_sim_and_wait(self, command, match_string, simulation): logger.debug("Match string found") yield p - os.killpg(os.getpgid(p.pid), signal.SIGTERM) - - @pytest.fixture() - def simulation(self, tmp_path): - old_cwd = Path.cwd() - os.chdir(tmp_path) - yield os.chdir(old_cwd) + os.killpg(os.getpgid(p.pid), signal.SIGTERM) @pytest.fixture() def client(self, run_sim_and_wait, universe): @@ -108,44 +104,21 @@ def test_compare_imd_to_true_traj(self, universe, client, first_frame): for ts in universe.trajectory[first_frame:]: imdf = client.get_imdframe() if imdsinfo.time: - assert_allclose(imdf.time, ts.time) + assert_allclose(imdf.time, ts.time, atol=1e-03) assert_allclose(imdf.step, ts.data["step"]) if imdsinfo.box: assert_allclose_with_logging( imdf.box, ts.triclinic_dimensions, - atol=1e-02, + atol=1e-03, ) if imdsinfo.positions: assert_allclose_with_logging( - imdf.positions, ts.positions, atol=1e-02 + imdf.positions, ts.positions, atol=1e-03 ) if imdsinfo.velocities: assert_allclose_with_logging( - imdf.velocities, ts.velocities, atol=1e-02 + imdf.velocities, ts.velocities, atol=1e-03 ) if imdsinfo.forces: - assert_allclose_with_logging(imdf.forces, ts.forces, atol=1e-02) - - # imdf = client.get_imdframe() - # if imdsinfo.time: - # assert_allclose(imdf.time, universe.trajectory[1].time) - # assert_allclose(imdf.step, universe.trajectory[1].data["step"]) - # if imdsinfo.box: - # assert_allclose_with_logging( - # imdf.box, - # universe.trajectory[1].triclinic_dimensions, - # atol=1e-03, - # ) - # if imdsinfo.positions: - # assert_allclose_with_logging( - # imdf.positions, universe.trajectory[1].positions, atol=1e-03 - # ) - # if imdsinfo.velocities: - # assert_allclose_with_logging( - # imdf.velocities, universe.trajectory[1].velocities, atol=1e-03 - # ) - # if imdsinfo.forces: - # assert_allclose_with_logging( - # imdf.forces, universe.trajectory[1].forces, atol=1e-03 - # ) + assert_allclose_with_logging(imdf.forces, ts.forces, atol=1e-03) diff --git a/imdclient/tests/test_vdos.py b/imdclient/tests/test_vdos.py new file mode 100644 index 0000000..9d330b7 --- /dev/null +++ b/imdclient/tests/test_vdos.py @@ -0,0 +1,145 @@ +import pytest +from imdclient.vdos import VDOS +import MDAnalysis as mda +from MDAnalysis.lib.log import ProgressBar + + +TOPOL = "/scratch/mheyden1/POPC/run-NVE.tpr" +TRAJ = "/scratch/mheyden1/POPC/run-NVE.trr" + +# Source of truth + +import numpy as np +from scipy.fft import fft + + +class vdos: + def __init__(self, sel, nCorr): + self.sel = sel + self.nCorr = nCorr + self.nRes = sel.residues.n_residues + # initialize lists and arrays + # Note: atMassLists are lists of numpy arrays + # => they are not a numpy arrays themselves + self.atMassLists = ( + [] + ) # list of numpy arrays of atomic masses for each residue + for res in self.sel.residues: + # Note: np.newaxis is used to make the array 2D + self.atMassLists.append( + res.atoms.masses[:, np.newaxis].astype("float64") + ) + # numpy array of residue masses from MDAnalysis + self.resMassList = sel.residues.masses + # numpy array buffers for COM position, COM velocity, angular momentum, and (sorted) moments of inertia + self.COMposBuffer = np.zeros((nCorr, self.nRes, 3), dtype=np.float64) + self.COMvelBuffer = np.zeros((nCorr, self.nRes, 3), dtype=np.float64) + # time and frequency axes for correlation functions and VDoS + self.tau = np.zeros(nCorr, dtype=np.float64) + self.wavenumber = np.zeros(nCorr, dtype=np.float64) + # numpy arrays for correlation functions and VDoS + # -> rigid body translation + self.trCorr = np.zeros((nCorr, self.nRes), dtype=np.float64) + self.trVDoS = np.zeros((nCorr, self.nRes), dtype=np.float64) + # initialize counter for normalization + self.corrCnt = np.zeros(self.nRes, dtype=int) + # flag for normalization + self.normalized = 0 + + def processStep(self, tStep, time): + """process a single time step of the simulation""" + idx = tStep % self.nCorr + if tStep < self.nCorr: + self.tau[tStep] = time + pos = [] + vel = [] + r = 0 + # compute center of mass position and velocity + for res in self.sel.residues: + pos.append(res.atoms.positions.astype("float64")) + vel.append(res.atoms.velocities.astype("float64")) + # self.COMposBuffer[idx,r] = res.atoms.center_of_mass() # => too slow & no equivalent for velocity + self.COMposBuffer[idx, r] = ( + np.sum(self.atMassLists[r] * pos[-1], axis=0) + / self.resMassList[r] + ) + self.COMvelBuffer[idx, r] = ( + np.sum(self.atMassLists[r] * vel[-1], axis=0) + / self.resMassList[r] + ) + r += 1 + # if sufficient data is available in buffers, compute correlation functions + if tStep >= self.nCorr - 1: + self.calcCorr(tStep + 1) + + def calcCorr(self, start): + """compute correlation functions for all data in buffers""" + # compute time correlation function for COM translation (for each residue) + for i in range(self.nCorr): + j = start % self.nCorr + k = (j + i) % self.nCorr + # for each residue + self.trCorr[i] += np.sum( + self.COMvelBuffer[j] * self.COMvelBuffer[k], axis=1 + ) + self.corrCnt += 1 + + def normalize(self): + """normalize correlation functions by number of data points""" + if self.normalized == 0: + for i in range(self.nRes): + self.trCorr[:, i] *= self.resMassList[i] / self.corrCnt[i] + self.normalized = 1 + + def calcVDOS(self): + """compute vibrational density of states from time correlation functions""" + period = (self.tau[1] - self.tau[0]) * (2 * self.nCorr - 1) + wn0 = (1.0 / period) * 33.35641 + self.wavenumber = np.arange(0, self.nCorr) * wn0 + tmp1 = np.zeros(2 * self.nCorr - 1, dtype=np.float64) + for i in range(self.nRes): + for j in range(self.nCorr): + tmp1[j] = self.trCorr[j][i] + for j in range(1, self.nCorr): + k = 2 * self.nCorr - j - 1 + tmp1[k] = tmp1[j] + tmp1 = fft(tmp1) + for j in range(self.nCorr): + self.trVDoS[j][i] = tmp1[j].real + + +@pytest.fixture +def true_vdos(): + u = mda.Universe(TOPOL, TRAJ) + sel = u.select_atoms("resname POPC") + true_vdos = vdos(sel, 200) + tStep = 0 + for ts in ProgressBar(u.trajectory[:400]): + true_vdos.processStep(tStep, ts.time) + tStep += 1 + true_vdos.normalize() + true_vdos.calcVDOS() + yield ( + true_vdos.trCorr, + true_vdos.trVDoS, + true_vdos.wavenumber, + true_vdos.tau, + ) + + +@pytest.fixture +def mda_vdos(): + u = mda.Universe(TOPOL, TRAJ) + sel = u.select_atoms("resname POPC") + vdos = VDOS(u.trajectory, sel, nCorr=200).run(stop=400) + yield ( + vdos.results.trCorr, + vdos.results.trVDoS, + vdos.results.wavenumber, + vdos.results.tau, + ) + + +def test_compare_vdos(true_vdos, mda_vdos): + for true, mda in zip(true_vdos, mda_vdos): + np.testing.assert_allclose(true, mda, rtol=1e-3) diff --git a/imdclient/vdos.py b/imdclient/vdos.py new file mode 100644 index 0000000..5f7511f --- /dev/null +++ b/imdclient/vdos.py @@ -0,0 +1,419 @@ +import MDAnalysis as mda +from scipy.fft import rfft, fft +import numpy as np +from MDAnalysis.analysis.base import ( + AnalysisBase, +) +from MDAnalysis.lib.log import ProgressBar + +import logging + +logger = logging.getLogger(__name__) + + +class StreamFriendlyAnalysisBase(AnalysisBase): + _analysis_algorithm_is_parallelizable = False + + def __init__(self, trajectory, verbose=False, **kwargs): + super().__init__(trajectory, verbose=verbose, **kwargs) + if hasattr(trajectory, "one_pass") and trajectory.one_pass: + self._one_pass = True + else: + self._one_pass = False + + def run( + self, + start=None, + stop=None, + step=None, + frames=None, + verbose=None, + *, + progressbar_kwargs={}, + ): + """Perform the calculation + + Parameters + ---------- + start : int, optional + start frame of analysis + stop : int, optional + stop frame of analysis + step : int, optional + number of frames to skip between each analysed frame + frames : array_like, optional + array of integers or booleans to slice trajectory; `frames` can + only be used *instead* of `start`, `stop`, and `step`. Setting + *both* `frames` and at least one of `start`, `stop`, `step` to a + non-default value will raise a :exc:`ValueError`. + + .. versionadded:: 2.2.0 + + verbose : bool, optional + Turn on verbosity + + progressbar_kwargs : dict, optional + ProgressBar keywords with custom parameters regarding progress bar position, etc; + see :class:`MDAnalysis.lib.log.ProgressBar` for full list. + + + .. versionchanged:: 2.2.0 + Added ability to analyze arbitrary frames by passing a list of + frame indices in the `frames` keyword argument. + + .. versionchanged:: 2.5.0 + Add `progressbar_kwargs` parameter, + allowing to modify description, position etc of tqdm progressbars + """ + logger.info("Choosing frames to analyze") + # if verbose unchanged, use class default + verbose = ( + getattr(self, "_verbose", False) if verbose is None else verbose + ) + + if self._one_pass: + self._sliced_trajectory = self._trajectory + self.frames = np.zeros(1000, dtype=int) + self.times = np.zeros(1000) + else: + self._setup_frames( + self._trajectory, + start=start, + stop=stop, + step=step, + frames=frames, + ) + + logger.info("Starting preparation") + self._prepare() + + if self._one_pass: + logger.info( + "Starting analysis loop over trajectory frames", + ) + else: + logger.info( + "Starting analysis loop over %d trajectory frames", + self.n_frames, + ) + + for i, ts in enumerate( + ProgressBar( + self._sliced_trajectory, verbose=verbose, **progressbar_kwargs + ) + ): + self._frame_index = i + self._ts = ts + if self._one_pass: + self._alloc_frame_and_time_arrays() + self.frames[i] = ts.frame + self.times[i] = ts.time + self._single_frame() + logger.info("Finishing up") + + if self._one_pass: + self.frames = self.frames[: self._frame_index + 1] + self.times = self.times[: self._frame_index + 1] + self.n_frames = self._frame_index + 1 + + self._conclude() + + return self + + def _alloc_frame_and_time_arrays(self): + if self._frame_index == len(self.frames) - 1: + self.frames = np.resize(self.frames, 1000 + len(self.frames)) + self.times = np.resize(self.times, 1000 + len(self.times)) + + +class VDOS(StreamFriendlyAnalysisBase): + + def __init__( + self, trajectory, selection, nCorr=100, verbose=False, **kwargs + ): + super().__init__(trajectory, verbose=False, **kwargs) + self.sel = selection + self.nCorr = nCorr + self.nRes = selection.residues.n_residues + + def _prepare(self): + # numpy array buffers for COM position, COM velocity, angular momentum, and (sorted) moments of inertia + self.COMposBuffer = np.zeros( + (self.nCorr, self.nRes, 3), dtype=np.float64 + ) + self.COMvelBuffer = np.zeros( + (self.nCorr, self.nRes, 3), dtype=np.float64 + ) + + # time and frequency axes for correlation functions and VDoS + self.results["tau"] = np.zeros(self.nCorr, dtype=np.float64) + self.results["wavenumber"] = np.zeros(self.nCorr, dtype=np.float64) + # numpy arrays for correlation functions and VDoS + # -> rigid body translation + self.results["trCorr"] = np.zeros( + (self.nCorr, self.nRes), dtype=np.float64 + ) + self.results["trVDoS"] = np.zeros( + (self.nCorr, self.nRes), dtype=np.float64 + ) + # initialize counter for normalization + self.corrCnt = 0 + + # # If residues can have different num atoms, use list + # self.atMassLists = [] + # for res in self.sel.residues: + # self.atMassLists.append( + # res.atoms.masses[:, np.newaxis].astype("float64") + # ) + + self.atMassLists = np.empty( + (len(self.sel.residues), len(self.sel.residues[0].atoms), 1), + dtype=np.float64, + ) + for i, res in enumerate(self.sel.residues): + self.atMassLists[i] = res.atoms.masses[:, np.newaxis].astype( + "float64" + ) + + def _single_frame(self): + idx = self._ts.frame % self.nCorr + + if self._ts.frame < self.nCorr: + self.results.tau[idx] = self._ts.time + + atMassLists = self.atMassLists + sel = self.sel + masses = sel.residues.masses + COMvelBuffer = self.COMvelBuffer + + # # compute center of mass position and velocity + for i in range(self.nRes): + vel = sel.residues[i].atoms.velocities.astype("float64") + COMvelBuffer[idx, i] = ( + np.sum(atMassLists[i] * vel, axis=0) / masses[i] + ) + + # if sufficient data is available in buffers, compute correlation functions + if self._ts.frame >= self.nCorr - 1: + self._calcCorr(self._ts.frame + 1) + + def _calcCorr(self, start): + """compute correlation functions for all data in buffers""" + # compute time correlation function for COM translation (for each residue) + # can be parallelized + trCorr = self.results.trCorr + COMvelBuffer = self.COMvelBuffer + nCorr = self.nCorr + + for i in range(nCorr): + j = start % nCorr + k = (j + i) % nCorr + trCorr[i] += np.sum(COMvelBuffer[j] * COMvelBuffer[k], axis=1) + self.corrCnt += 1 + + def _conclude(self): + """normalize correlation functions by number of data points + and calculate compute vibrational density of states from time correlation functions + """ + # Normalization + # multiply by mass to convert velocity -> momentum + self.results.trCorr[:] *= self.sel.residues.masses[:] / self.corrCnt + ## Calculate VDoS + period = (self.results.tau[1] - self.results.tau[0]) * ( + 2 * self.nCorr - 1 + ) + wn0 = (1.0 / period) * 33.35641 + self.results.wavenumber = np.arange(0, self.nCorr) * wn0 + tmp1 = np.zeros((2 * self.nCorr - 1, self.nRes), dtype=np.float64) + tmp1[: self.nCorr] = self.results.trCorr[:] + tmp1[self.nCorr :] = tmp1[1 : self.nCorr][::-1] + self.results.trVDoS = rfft(tmp1, axis=0) + + +class roVDoS(StreamFriendlyAnalysisBase): + + def __init__( + self, trajectory, selection, nCorr=100, verbose=False, **kwargs + ): + super().__init__(trajectory, verbose=False, **kwargs) + self.sel = selection + self.nCorr = nCorr + self.nRes = selection.residues.n_residues + + def _prepare(self): + + self.AngMomentumBuffer = np.zeros( + (self.nCorr, self.nRes, 3), dtype=np.float64 + ) + self.COMposBuffer = np.zeros( + (self.nCorr, self.nRes, 3), dtype=np.float64 + ) + self.COMvelBuffer = np.zeros( + (self.nCorr, self.nRes, 3), dtype=np.float64 + ) + self.MOIBuffer = np.zeros((self.nCorr, self.nRes, 3)) + + # time and frequency axes for correlation functions and VDoS + self.results["tau"] = np.zeros(self.nCorr, dtype=np.float64) + self.results["wavenumber"] = np.zeros(self.nCorr, dtype=np.float64) + + self.results["roCorr"] = np.zeros( + (self.nCorr, self.nRes), dtype=np.float64 + ) + self.results["roVDoS"] = np.zeros( + (self.nCorr, self.nRes), dtype=np.float64 + ) + + self.corrCnt = 0 + + # If residues can have different num atoms, use list + self.atMassLists = [] + for res in self.sel.residues: + self.atMassLists.append( + res.atoms.masses[:, np.newaxis].astype("float64") + ) + + self.prev_evecs = None + + def _single_frame(self): + idx = self._ts.frame % self.nCorr + + if self._ts.frame < self.nCorr: + self.results.tau[idx] = self._ts.time + + atMassLists = self.atMassLists + sel = self.sel + residue_masses = sel.residues.masses + COMvelBuffer = self.COMvelBuffer + COMposBuffer = self.COMposBuffer + AngMomentumBuffer = self.AngMomentumBuffer + MOIBuffer = self.MOIBuffer + prev_evecs = self.prev_evecs + + calc_inertia_tensor = self._calc_inertia_tensor + + for i in range(self.nRes): + atoms = sel.residues[i].atoms + atom_masses_col = atMassLists[i] + atom_masses_row = atoms.masses + residue_mass = residue_masses[i] + + ## Compute COM pos + pos = atoms.positions.astype("float64") + com_pos = np.sum(atom_masses_col * pos, axis=0) / residue_mass + COMposBuffer[idx, i] = com_pos + + # valid_com = sel.residues[i].atoms.center_of_mass() + # if not np.allclose(com_pos, valid_com): + # print(f"Center of mass is not valid {com_pos} != {valid_com}") + + ## Compute COM pos + vel = atoms.velocities.astype("float64") + com_vel = np.sum(atom_masses_col * vel, axis=0) / residue_mass + COMvelBuffer[idx, i] = com_vel + + ## Calcular angular momentum + ang = np.sum( + atom_masses_col * np.cross((pos - com_pos), (vel - com_vel)), + axis=0, + ) + + ## Compute MOI, Angular Momentum along principal axis + inertia_tensor = calc_inertia_tensor(com_pos, pos, atom_masses_row) + # valid_inertia_tensor = sel.residues[i].atoms.moment_of_inertia() + # if not np.allclose(inertia_tensor, valid_inertia_tensor): + # print( + # f"Inertia tensor is not valid {inertia_tensor} != {valid_inertia_tensor}" + # ) + + values, evecs = np.linalg.eigh(inertia_tensor) + + ## Ensure that the eigenvectors don't flip + if self.prev_evecs is not None: + for j in range(3): + if np.dot(evecs[:, j], prev_evecs[:, j]) < 0: + evecs[:, j] *= -1 + prev_evecs = evecs + + indices = np.argsort(values) + rot_axis = evecs[:, indices] + + # valid_principal_axis = sel.residues[i].atoms.principal_axes().T + # if not np.allclose(rot_axis, valid_principal_axis): + # print( + # f"Principal axis is not valid {rot_axis} != {valid_principal_axis}" + # ) + + MOIBuffer[idx, i] = values + + AngMomentumBuffer[idx, i] = np.dot(ang, rot_axis) + # np.sum(ang * rot_axis, axis=0) + + # if sufficient data is available in buffers, compute correlation functions + if self._ts.frame >= self.nCorr - 1: + self._calcCorr(self._ts.frame + 1) + + def _calc_inertia_tensor(self, com, pos, masses): + pos_com_frame = pos - com + tens = np.zeros((3, 3), dtype=np.float64) + # xx + tens[0][0] = ( + masses * (pos_com_frame[:, 1] ** 2 + pos_com_frame[:, 2] ** 2) + ).sum() + # xy & yx + tens[0][1] = tens[1][0] = -( + masses * pos_com_frame[:, 0] * pos_com_frame[:, 1] + ).sum() + # xz & zx + tens[0][2] = tens[2][0] = -( + masses * pos_com_frame[:, 0] * pos_com_frame[:, 2] + ).sum() + # yy + tens[1][1] = ( + masses * (pos_com_frame[:, 0] ** 2 + pos_com_frame[:, 2] ** 2) + ).sum() + # yz + zy + tens[1][2] = tens[2][1] = -( + masses * pos_com_frame[:, 1] * pos_com_frame[:, 2] + ).sum() + # zz + tens[2][2] = ( + masses * (pos_com_frame[:, 0] ** 2 + pos_com_frame[:, 1] ** 2) + ).sum() + return tens + + def _calcCorr(self, start): + """compute correlation functions for all data in buffers""" + # compute time correlation function for COM translation (for each residue) + # can be parallelized + roCorr = self.results.roCorr + MOIBuffer = self.MOIBuffer + nCorr = self.nCorr + AngMomentumBuffer = self.AngMomentumBuffer + + for i in range(nCorr): + j = start % nCorr + k = (j + i) % nCorr + roCorr[i] += np.sum( + (AngMomentumBuffer[j] / np.sqrt(MOIBuffer[j])) + * (AngMomentumBuffer[k] / np.sqrt(MOIBuffer[k])), + axis=1, + ) + self.corrCnt += 1 + + def _conclude(self): + """normalize correlation functions by number of data points + and calculate compute vibrational density of states from time correlation functions + """ + # Normalization + self.results.roCorr[:] /= self.corrCnt + ## Calculate VDoS + period = (self.results.tau[1] - self.results.tau[0]) * ( + 2 * self.nCorr - 1 + ) + wn0 = (1.0 / period) * 33.35641 + self.results.wavenumber = np.arange(0, self.nCorr) * wn0 + tmp1 = np.zeros((2 * self.nCorr - 1, self.nRes), dtype=np.float64) + tmp1[: self.nCorr] = self.results.roCorr[:] + tmp1[self.nCorr :] = tmp1[1 : self.nCorr][::-1] + self.results.roVDoS = rfft(tmp1, axis=0) diff --git a/pyproject.toml b/pyproject.toml index 73a46fa..ff7ef84 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,7 +32,7 @@ test = [ "pytest>=6.0", "pytest-xdist>=2.5", "pytest-cov>=3.0", - "MDAnalysis>=2.7.0", + "MDAnalysisTests>=2.7.0", ] doc = [ "sphinx",