Repository: jmborr/nscsim

Issue:
coherent intermediate dynamic scattering function

Discussion:

Patch:
diff --git a/.travis.yml b/.travis.yml
index ab323bd..23e7657 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -1,6 +1,7 @@
 language: python
 
 python:
+  - "2.7"
   - "3.4"
   - "3.6"
 
diff --git a/notebooks/quasielastic/__init__.py b/notebooks/quasielastic/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/notebooks/quasielastic/coherent.ipynb b/notebooks/quasielastic/coherent.ipynb
new file mode 100644
index 0000000..154097b
--- /dev/null
+++ b/notebooks/quasielastic/coherent.ipynb
@@ -0,0 +1,286 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Requirements\n",
+    "\n",
+    "#### Python packages\n",
+    "\n",
+    "The following packages are not required by <i>nscsim</i> but are needed to run this notebook\n",
+    "\n",
+    "- [cython](https://pypi.org/project/Cython/)\n",
+    "- [MDAnalysis](https://pypi.org/project/MDAnalysis/)\n",
+    "- [periodictable](https://pypi.org/project/periodictable/)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Requirements"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 30,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%matplotlib notebook\n",
+    "import matplotlib.pyplot as plt\n",
+    "plt.ion()\n",
+    "import numpy as np\n",
+    "import MDAnalysis as mda\n",
+    "import periodictable\n",
+    "import nscsim\n",
+    "from nscsim.quasielastic import coherent as coh\n",
+    "from nscsim import plotting"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Download water simulation\n",
+    "`wget` and `tar` commands are necessary.  \n",
+    "You may skip this step if running the notebook more than once"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%capture\n",
+    "%%bash\n",
+    "wget https://www.dropbox.com/s/7imfnz6ma1kh41w/water_4000.tar.gz\n",
+    "tar zxf water_4000.tar.gz"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Create system and extract coordinates\n",
+    "\n",
+    "`u` is a `MDAnalysis.Universe` object. It contains 154 Lithium atoms, 154 Chloride atoms, and 1128 TIP3 water molecules."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 31,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "u = mda.Universe('topology.psf', 'trajectory.dcd')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We investigate the trajectory."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 32,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "number of conformations = 4000\n",
+      "Size of the box, in Angstroms = [33.6035 33.6035 33.6035]\n",
+      "It is a cubic box, as evidenced by the angles between box sides:  [90. 90. 90.]\n"
+     ]
+    }
+   ],
+   "source": [
+    "print('number of conformations =', u.trajectory.n_frames)\n",
+    "print('Size of the box, in Angstroms =', u.dimensions[0:3])\n",
+    "print('It is a cubic box, as evidenced by the angles between box sides: ',u.dimensions[3:])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Select only the hydrogens and oxygens. There should be three atoms per water molecule. A look at file <i>topology.psf</i> shows that water hydrogens have type `HT` and water oxygens have type `OT`"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 33,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "number of atoms selected = 3384\n",
+      "number of atoms per water molecule =  3.0\n"
+     ]
+    }
+   ],
+   "source": [
+    "sel = u.select_atoms('type HT or type OT')\n",
+    "print('number of atoms selected =', len(sel))\n",
+    "print('number of atoms per water molecule = ', len(sel)/len(sel.residues))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Obtain the bound coherent neutron scattering lengths for the atoms of our selection `sel`. Notice we have to \"translate\" from NAMD atom types to the name of the atomic element. We use dictionary `type_to_element` for this. Then we iterate over our selected atoms"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 34,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "type_to_element = dict(HT='H', OT='O')\n",
+    "#\n",
+    "# We use a for loop here\n",
+    "#\n",
+    "b_c = list()\n",
+    "for atom in sel:\n",
+    "    atomic_element = type_to_element[atom.type]\n",
+    "    coherent_length = getattr(periodictable, atomic_element).neutron.b_c\n",
+    "    b_c.append(coherent_length)\n",
+    "b_c = np.asarray(b_c)\n",
+    "#\n",
+    "# Advanced python: the previous lines can be substituted with the following one-liner\n",
+    "#\n",
+    "#b_c = np.asarray([getattr(periodictable, type_to_element[at.type]).neutron.b_c for at in sel])\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Calculate the coherent dynamic structure factor `S(q)`\n",
+    "\n",
+    "We calculate scattering only for Q=0.5 using the first 1000 frames in the trajectory. The liquid allows to compute scattering only as a function of Q-vector modulus so we carry out a spherical average for all orientations of vector Q (function `intermediate_spherical`).\n",
+    "\n",
+    "Notice we call `intermediate_spherical` only with one core (`n_cores=1`) because in most computers numpy will automatically use parallel versions of underlying computing BLAS libraries. Thus, all cores will be used even if we require one core. If this is not the case for your computer, then you can request more cores by changing the value of keyword `n_cores`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tr = np.asarray([sel.positions for _ in u.trajectory[0:1000:1]])  # extract coordinates to numpy array\n",
+    "q_mod = np.array([0.5, ])\n",
+    "sf = coh.intermediate_spherical(tr, q_mod, b_c, n_cores=1)\n",
+    "times = coh.times(tr)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Simple plotting of the curve."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plotting.single_curve(times, sf[0], label='time', ylabel='Coherent Intensity')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Notice that at long times the scattering is not reliable due to low statistical sampling. A more reliable curve at times~1000 can be obtained if we use all frames, not just the first 1000:\n",
+    "\n",
+    "<code>tr = np.asarray([sel.positions for _ in u.trajectory[::1]]) # all frames in the trajectory</code>"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Calculate the coherent dynamic structure factor `S(q)` in a log scale of Q-values\n",
+    "\n",
+    "We will calculate more values of Q using the first 100 oxygen atoms and all the trajectory frames. Here we calculate scattering from `Q=0.01` to `Q=10` for a total of 15 Q-values spaced in a log scale."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sel = u.select_atoms('type OT')[0:100]\n",
+    "tr = np.asarray([sel.positions for _ in u.trajectory[::1]])  # extract coordinates to numpy array\n",
+    "b_c = np.asarray([getattr(periodictable, 'O').neutron.b_c for at in sel])\n",
+    "\n",
+    "q_mod = nscsim.qvec.moduli_logscale(min_exp=-2, max_exp=1, n_per_base=5, base=10)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sf = coh.intermediate_spherical(tr, q_mod, b_c, n_cores=1)\n",
+    "times = coh.times(tr)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Plot all curves sequentially"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "common = dict(xlabel='time', ylabel='Coherent Intensity', xscale='linear', yscale='linear')\n",
+    "for i, curve in enumerate(sf):\n",
+    "    common['label'] = 'Q={:.2f}'.format(q_mod[i])\n",
+    "    plotting.single_curve(times, curve, **common)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.6.7"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/nscsim/plotting.py b/nscsim/plotting.py
new file mode 100644
index 0000000..33d5516
--- /dev/null
+++ b/nscsim/plotting.py
@@ -0,0 +1,28 @@
+import matplotlib.pyplot as plt
+
+
+def single_curve(x, y, xlabel='x', ylabel='y', xscale='linear',
+                 yscale='linear', **plotkwargs):
+    r"""
+    matplotlib.axes.Axes.plot figure
+
+    Parameters
+    ----------
+    x: numpy.ndarray
+        Abcisa values
+    y: numpy.ndarray
+        Ordinate values
+    plotkwargs: dict
+        Additional arguments to matplotlib.axes.Axes.plot()
+    """
+    if plotkwargs is None:
+        plotkwargs = dict()
+    fig, ax = plt.subplots(1, 1)
+    ax.plot(x, y, **plotkwargs)
+    ax.set_xlabel(xlabel, size=25)
+    ax.set_ylabel(ylabel, size=25)
+    ax.set_yscale(yscale)
+    ax.set_xscale(xscale)
+    ax.legend()
+    plt.tight_layout()
+    plt.show()
diff --git a/nscsim/quasielastic/coherent.py b/nscsim/quasielastic/coherent.py
new file mode 100644
index 0000000..bf47aff
--- /dev/null
+++ b/nscsim/quasielastic/coherent.py
@@ -0,0 +1,171 @@
+from __future__ import (print_function, absolute_import)
+
+import numpy as np
+from nscsim.utilities import (serial_bar, parallel_bar, glog)
+from nscsim import qvec
+import multiprocessing
+import pathos
+
+
+def intermediate_amplitudes(tr, q, bc, n_cores=None):
+    r"""
+    For every frame, calculate the coherent amplitude.
+
+    The coherent amplitude is a complex number
+    math:`A(\vec{q}) = sum_i b_i e^{-i\vec{q}\vec(r)_i}
+
+    Parameters
+    ----------
+    tr: numpy.ndarray
+        Atomic trajectory, shape=(#frames, #atoms, 3)
+    q: numpy.ndarray
+        Array of q-vectors shape=(#vectors, 3)
+    bc: numpy.ndarray
+        Coherent scattering lengths, shape=(#atoms, 1)
+    n_cores: int
+        Number of CPU's to use. All cores minus 1, or 1 if only 1 core
+
+    Returns
+    -------
+    np.ndarray
+        Scattering amplitudes shape=(#frames, #q's)
+    """
+
+    def serial_worker(frame):
+        r"""
+        Calculate q-dependent coherent amplitudes for one frame
+
+        Parameters
+        ----------
+        frame: numpy.ndarray
+            shape = (#atoms, 3) coordinates of the system
+        """
+        exponents = np.tensordot(frame, q, axes=(1, 1))  # shape=(#atoms, #q's)
+        exponentials = np.exp(1j * exponents)
+        return np.tensordot(bc, exponentials, axes=1)  # shape=(#q's,)
+
+    glog.info('\nCalculating coherent amplitudes for one set of q vectors\n')
+    if n_cores is None:
+        n_cores = max(1, multiprocessing.cpu_count() - 1)
+    if n_cores == 1:
+        amps = np.asarray([serial_worker(f) for f in serial_bar(tr)])
+        return amps
+    else:
+        pool = pathos.pools.ProcessPool(ncpus=n_cores)
+        amps = np.asarray(parallel_bar(serial_worker, tr))
+        pool.terminate()
+        return amps
+
+
+def times(tr):
+    r"""
+    Times for the coherent scattering functions. Units are whatever the
+    time step between frames in the input trajectory.
+
+    Minimum time is `1-len(tr)` and maximum time is `len(tr)-1`
+
+    Parameters
+    ----------
+    tr: numpy.ndarray
+        Atomic trajectory, shape=(#frames, #atoms, 3)ndarray
+
+    Returns
+    -------
+    numpy.ndarray
+    """
+    return np.arange(1 - len(tr), len(tr))
+
+
+def intermediate_vector_set(tr, q, bc, n_cores=None):
+    r"""
+    Coherent scattering for a collection of vectors
+
+    Coherent scattering is reported for each vector in the set.
+    The scattering curve is symmetrized,
+    e.g. if the trajectory has 10 frames
+    then the scattering curve would start at t=-9
+    and runs up to t=9.
+
+
+    Parameters
+    ----------
+    tr: numpy.ndarray
+        Atomic trajectory, shape=(#frames, #atoms, 3)
+    q: numpy.ndarray
+        Array of q-vectors shape=(#vectors, 3)
+    bc: numpy.ndarray
+        Coherent scattering lengths, shape=(#atoms, 1)
+    n_cores: int
+        Number of CPU's to use. All cores minus 1, or 1 if only 1 core
+
+    Returns
+    -------
+    np.ndarray
+        shape = (len(q), 2 * #frames - 1)
+    """
+    amps = intermediate_amplitudes(tr, q, bc, n_cores=n_cores).transpose()
+    w = np.arange(1, 1 + len(tr))
+    w = 1.0 / np.concatenate((w, w[::-1][1:]))
+    sf = [w * np.correlate(a, a, 'full') for a in amps]
+    return np.asarray(sf)
+
+
+def intermediate(tr, q, bc, n_cores=None,
+                 averaging=(np.average, None, dict(axis=0))):
+    r"""
+    Coherent scattering averaged over scattering for a collection of q-vectors
+
+    Parameters
+    ----------
+    tr: numpy.ndarray
+        Atomic trajectory, shape=(#frames, #atoms, 3)
+    q: numpy.ndarray
+        Array of q-vectors shape=(#vectors, 3)
+    bc: numpy.ndarray
+        Coherent scattering lengths, shape=(#atoms, 1)
+    n_cores: int
+        Number of CPU's to use. All cores minus 1, or 1 if only 1 core
+    averaging: tuple
+        The Q-average to carry out over the :math:`S(\vec{q}, t)`, the
+        return value from `intermediate_vector_set`.
+        Tuple elements:
+            - averaging function
+            - positional arguments for function, excluded the first which is
+              assumed to be array :math:`S(\vec{q}, t)`
+            - keyword arguments for function
+
+    Returns
+    -------
+    np.ndarray
+        shape = (#frames)
+    """
+    sf = intermediate_vector_set(tr, q, bc, n_cores=n_cores)
+    av_args = list() if averaging[1] is None else averaging[1]
+    return averaging[0](sf, *av_args, **averaging[2])
+
+
+def intermediate_spherical(tr, q_mod, b_c, n_cores=1):
+    r"""
+    Coherent scattering averaged over scattering for a collection of q-vectors
+
+    Parameters
+    ----------
+    tr: numpy.ndarray
+        Atomic trajectory, shape=(#frames, #atoms, 3)
+    q_mod: numpy.ndarray
+        Array of q-vector moduli shape=(#q-moduli, 1)
+    b_c: numpy.ndarray
+        Coherent scattering lengths, shape=(#atoms, 1)
+    n_cores: int
+        Number of CPU's to use. All cores minus 1, or 1 if only 1 core
+
+    Returns
+    -------
+    np.ndarray
+        shape = (#frames)
+    """
+    q_sets = qvec.sphere_average(q_mod)
+    glog.info('\nCalculating coherent scattering for a set of q-moduli\n')
+    sf = [np.real(intermediate(tr, q_set, b_c, n_cores=1))
+          for q_set in serial_bar(q_sets)]
+    return np.asarray(sf)
diff --git a/nscsim/qvec.py b/nscsim/qvec.py
index 73a198e..5427029 100644
--- a/nscsim/qvec.py
+++ b/nscsim/qvec.py
@@ -35,9 +35,9 @@ def moduli_linscale(q_mod_min, q_mod_max=None, q_mod_delta=None,
     r"""
     An array of q vector moduli (q_mods) is generated depending on the passed
     arguments. These are the options:
-    1. (q_mod_min, q_mod_max, q_mod_delta) generates q_mods
-    2. (q_mod_min, q_mod_max, n_q_mod) generates q_mods
-    3. (q_mod_min, q_mod_delta, n_q_mod) generates q_mods
+    1. (q_mod_min, q_mod_max, q_mod_delta)
+    2. (q_mod_min, q_mod_max, n_q_mod) with guaranteed q_mod_min and q_mod_max
+    3. (q_mod_min, q_mod_delta, n_q_mod)
 
     Parameters
     ----------
@@ -61,9 +61,8 @@ def moduli_linscale(q_mod_min, q_mod_max=None, q_mod_delta=None,
             # q_mods from triad (q_mod_min, q_mod_max, q_mod_delta)
             q_mods = np.arange(q_mod_min, q_mod_max, q_mod_delta)
         elif n_q_mod is not None:
-            # q_mods from triad (q_mod_min, q_mod_max, n_q_mod)
-            q_mod_delta = (q_mod_max - q_mod_min) / n_q_mod
-            q_mods = np.arange(q_mod_min, q_mod_max, q_mod_delta)
+            q_mods = np.linspace(q_mod_min, q_mod_max, num=n_q_mod,
+                                 endpoint=True)
     elif q_mod_delta is not None and n_q_mod is not None:
         # q_mods from triad (q_mod_min, q_mod_delta, n_q_mod)
         q_mods = q_mod_min + q_mod_delta * np.arange(n_q_mod)
diff --git a/nscsim/utilities.py b/nscsim/utilities.py
index 69bb2a9..2c33d09 100644
--- a/nscsim/utilities.py
+++ b/nscsim/utilities.py
@@ -1,10 +1,25 @@
 from __future__ import (absolute_import, division, print_function)
 
+import logging
 import functools
 import numpy as np
 from collections import namedtuple, Mapping
 import multiprocessing
 import ctypes
+import p_tqdm
+
+
+glog = logging.getLogger('nscsim')
+glog.addHandler(logging.StreamHandler())
+glog.setLevel(logging.INFO)
+
+
+def serial_bar(iterable):
+    return p_tqdm.tqdm(iterable, miniters=int(len(iterable/100)))
+
+
+def parallel_bar(iterable):
+    return p_tqdm.p_map(iterable, miniters=int(len(iterable/100)))
 
 
 def namedtuplefy(func):
diff --git a/requirements.txt b/requirements.txt
index f6ce8e9..b0fa3a4 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -3,4 +3,5 @@ numpy
 pathos
 scipy
 six
-tqdm
+matplotlib
+p_tqdm


Test patch:
diff --git a/tests/unit/test_qvec.py b/tests/unit/test_qvec.py
index 5638a5a..3351649 100644
--- a/tests/unit/test_qvec.py
+++ b/tests/unit/test_qvec.py
@@ -4,7 +4,6 @@ import numpy as np
 from numpy.testing import assert_almost_equal, assert_array_equal
 from math import sqrt
 import pytest
-
 from nscsim import qvec
 
 
@@ -16,9 +15,10 @@ def test_unit_sphere_surf():
 
 def test_moduli_linscale():
     for qs in (qvec.moduli_linscale(0, q_mod_max=1.0, q_mod_delta=0.1),
-               qvec.moduli_linscale(0, q_mod_max=1.0, n_q_mod=10),
                qvec.moduli_linscale(0, q_mod_delta=0.1, n_q_mod=10)):
         assert np.array_equal(qs, np.arange(0, 1.0, 0.1))
+    qs = qvec.moduli_linscale(0, q_mod_max=1.0, n_q_mod=10)
+    assert_almost_equal(qs, np.linspace(0, 1, 10), decimal=2)
 
 
 def test_moduli_logscale():