jaxdecomp proto

.pre-commit-config.yaml

@@ -14,4 +14,4 @@ repos:
     rev: 5.13.2
     hooks:
       - id: isort
-        name: isort (python)
+        name: isort (python)

dev/jaxdecomp.py

@@ -0,0 +1,69 @@
+import argparse
+
+import jax
+import numpy as np
+
+# Setting up distributed jax
+jax.distributed.initialize()
+rank = jax.process_index()
+size = jax.process_count()
+
+import jax.numpy as jnp
+import jax_cosmo as jc
+from jax.experimental import mesh_utils
+from jax.sharding import Mesh
+
+from jaxpm.painting import cic_paint
+from jaxpm.pm import linear_field, lpt
+
+mesh_shape = [256, 256, 256]
+box_size = [256., 256., 256.]
+snapshots = jnp.linspace(0.1, 1., 2)
+
+
+@jax.jit
+def run_simulation(omega_c, sigma8, seed):
+    # Create a cosmology
+    cosmo = jc.Planck15(Omega_c=omega_c, sigma8=sigma8)
+
+    # Create a small function to generate the matter power spectrum
+    k = jnp.logspace(-4, 1, 128)
+    pk = jc.power.linear_matter_power(
+        jc.Planck15(Omega_c=omega_c, sigma8=sigma8), k)
+    pk_fn = lambda x: jc.scipy.interpolate.interp(x.reshape([-1]), k, pk
+                                                  ).reshape(x.shape)
+
+    # Create initial conditions
+    initial_conditions = linear_field(mesh_shape, box_size, pk_fn, seed=seed)
+
+    # Initialize particle displacements
+    dx, p, f = lpt(cosmo, initial_conditions, 1.0)
+
+    field = cic_paint(jnp.zeros_like(initial_conditions), dx)
+    return field
+
+
+def main(args):
+    # Setting up distributed random numbers
+    master_key = jax.random.PRNGKey(42)
+    key = jax.random.split(master_key, size)[rank]
+
+    # Create computing mesh and sharding information
+    devices = mesh_utils.create_device_mesh((2, 2))
+    mesh = Mesh(devices.T, axis_names=('x', 'y'))
+
+    # Run the simulation on the compute mesh
+    with mesh:
+        field = run_simulation(0.32, 0.8, key)
+
+    print('done')
+    np.save(f'field_{rank}.npy', field.addressable_data(0))
+
+    # Closing distributed jax
+    jax.distributed.shutdown()
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser("Distributed LPT N-body simulation.")
+    args = parser.parse_args()
+    main(args)

jaxpm/distributed.py

@@ -0,0 +1,151 @@
+from typing import Any, Callable, Hashable
+
+Specs = Any
+AxisName = Hashable
+
+try:
+    import jaxdecomp
+    distributed = True
+except ImportError:
+    print("jaxdecomp not installed. Distributed functions will not work.")
+    distributed = False
+
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+from jax._src import mesh as mesh_lib
+from jax.experimental.shard_map import shard_map
+from jax.sharding import PartitionSpec as P
+
+# NOTE
+# This should not be used as a decorator
+# Must be used inside a function only
+# Example
+# BAD
+# @autoshmap
+# def foo():
+#     pass
+# GOOD
+# def foo():
+#     return autoshmap(foo_impl)()
+
+
+def autoshmap(f: Callable,
+              in_specs: Specs,
+              out_specs: Specs,
+              check_rep: bool = True,
+              auto: frozenset[AxisName] = frozenset()):
+    """Helper function to wrap the provided function in a shard map if
+    the code is being executed in a mesh context."""
+    mesh = mesh_lib.thread_resources.env.physical_mesh
+    if mesh.empty:
+        return f
+    else:
+        return shard_map(f, mesh, in_specs, out_specs, check_rep, auto)
+
+
+def fft3d(x):
+    if distributed and not (mesh_lib.thread_resources.env.physical_mesh.empty):
+        return jaxdecomp.pfft3d(x.astype(jnp.complex64))
+    else:
+        return jnp.fft.fftn(x.astype(jnp.complex64))
+
+
+def ifft3d(x):
+    if distributed and not (mesh_lib.thread_resources.env.physical_mesh.empty):
+        return jaxdecomp.pifft3d(x).real
+    else:
+        return jnp.fft.ifftn(x).real
+
+
+def get_halo_size(halo_size):
+    mesh = mesh_lib.thread_resources.env.physical_mesh
+    if mesh.empty:
+        zero_ext = (0, 0, 0)
+        zero_tuple = (0, 0)
+        return (zero_tuple, zero_tuple, zero_tuple), zero_ext
+    else:
+        pdims = mesh.devices.shape
+    halo_x = (0, 0) if pdims[0] == 1 else (halo_size, halo_size)
+    halo_y = (0, 0) if pdims[1] == 1 else (halo_size, halo_size)
+
+    halo_x_ext = 0 if pdims[0] == 1 else halo_size // 2
+    halo_y_ext = 0 if pdims[1] == 1 else halo_size // 2
+    return ((halo_x, halo_y, (0, 0)), (halo_x_ext, halo_y_ext, 0))
+
+
+def halo_exchange(x, halo_extents, halo_periods=(True, True, True)):
+    mesh = mesh_lib.thread_resources.env.physical_mesh
+    if distributed and not (mesh.empty) and (halo_extents[0] > 0
+                                             or halo_extents[1] > 0):
+        return jaxdecomp.halo_exchange(x, halo_extents, halo_periods)
+    else:
+        return x
+
+
+def slice_unpad_impl(x, pad_width):
+
+    halo_x, _ = pad_width[0]
+    halo_y, _ = pad_width[0]
+
+    # Apply corrections along x
+    x = x.at[halo_x:halo_x + halo_x // 2].add(x[:halo_x // 2])
+    x = x.at[-(halo_x + halo_x // 2):-halo_x].add(x[-halo_x // 2:])
+    # Apply corrections along y
+    x = x.at[:, halo_y:halo_y + halo_y // 2].add(x[:, :halo_y // 2])
+    x = x.at[:, -(halo_y + halo_y // 2):-halo_y].add(x[:, -halo_y // 2:])
+
+    return x[halo_x:-halo_x, halo_y:-halo_y, :]
+
+
+def slice_pad(x, pad_width):
+    mesh = mesh_lib.thread_resources.env.physical_mesh
+    if distributed and not (mesh.empty) and (pad_width[0][0] > 0
+                                             or pad_width[1][0] > 0):
+        return autoshmap((partial(jnp.pad, pad_width=pad_width)),
+                         in_specs=(P('x', 'y')),
+                         out_specs=P('x', 'y'))(x)
+    else:
+        return x
+
+
+def slice_unpad(x, pad_width):
+    mesh = mesh_lib.thread_resources.env.physical_mesh
+    if distributed and not (mesh.empty) and (pad_width[0][0] > 0
+                                             or pad_width[1][0] > 0):
+        return autoshmap(partial(slice_unpad_impl, pad_width=pad_width),
+                         in_specs=(P('x', 'y')),
+                         out_specs=P('x', 'y'))(x)
+    else:
+        return x
+
+
+def get_local_shape(mesh_shape):
+    """ Helper function to get the local size of a mesh given the global size.
+  """
+    if mesh_lib.thread_resources.env.physical_mesh.empty:
+        return mesh_shape
+    else:
+        pdims = mesh_lib.thread_resources.env.physical_mesh.devices.shape
+        return [
+            mesh_shape[0] // pdims[0], mesh_shape[1] // pdims[1], mesh_shape[2]
+        ]
+
+
+def normal_field(mesh_shape, seed=None):
+    """Generate a Gaussian random field with the given power spectrum."""
+    if distributed and not (mesh_lib.thread_resources.env.physical_mesh.empty):
+        local_mesh_shape = get_local_shape(mesh_shape)
+        if seed is None:
+            key = None
+        else:
+            size = jax.process_count()
+            rank = jax.process_index()
+            key = jax.random.split(seed, size)[rank]
+        return autoshmap(
+            partial(jax.random.normal, shape=local_mesh_shape, dtype='float32'),
+            in_specs=P(None),
+            out_specs=P('x', 'y'))(key)  # yapf: disable
+    else:
+        return jax.random.normal(shape=mesh_shape, key=seed)

jaxpm/kernels.py

@@ -1,23 +1,49 @@
+from functools import partial
+
 import jax.numpy as jnp
+import jax_cosmo as jc
 import numpy as np
+from jax._src import mesh as mesh_lib
+from jax.sharding import PartitionSpec as P
+
+from jaxpm.distributed import autoshmap
 
 
-def fftk(shape, symmetric=True, finite=False, dtype=np.float32):
-    """ Return k_vector given a shape (nc, nc, nc) and box_size
+def fftk(shape, dtype=np.float32):
+    """
+    Generate Fourier transform wave numbers for a given mesh.
+
+    Args:
+        nc (int): Shape of the mesh grid.
+
+    Returns:
+        list: List of wave number arrays for each dimension in
+        the order [kx, ky, kz].
   """
-    k = []
-    for d in range(len(shape)):
-        kd = np.fft.fftfreq(shape[d])
-        kd *= 2 * np.pi
-        kdshape = np.ones(len(shape), dtype='int')
-        if symmetric and d == len(shape) - 1:
-            kd = kd[:shape[d] // 2 + 1]
-        kdshape[d] = len(kd)
-        kd = kd.reshape(kdshape)
+    kx, ky, kz = [jnp.fft.fftfreq(s, dtype=dtype) * 2 * np.pi for s in shape]
+
+    @partial(autoshmap,
+             in_specs=(P('x'), P('y'), P(None)),
+             out_specs=(P('x'), P(None, 'y'), P(None)))
+    def get_kvec(ky, kz, kx):
+        return (ky.reshape([-1, 1, 1]),
+                kz.reshape([1, -1, 1]),
+                kx.reshape([1, 1, -1])) # yapf: disable
+
+    if not mesh_lib.thread_resources.env.physical_mesh.empty:
+        ky, kz, kx = get_kvec(ky, kz, kx)  # The order corresponds
+    else:
+        kx, ky, kz = get_kvec(kx, ky, kz)  # The order corresponds
+
+    # to the order of dimensions in the transposed FFT
+    return kx, ky, kz
+
+
+def interpolate_power_spectrum(input, k, pk):
 
-        k.append(kd.astype(dtype))
-    del kd, kdshape
-    return k
+    pk_fn = lambda x: jc.scipy.interpolate.interp(x.reshape(-1), k, pk
+                                                  ).reshape(x.shape)
+    return autoshmap(pk_fn, in_specs=P('x', 'y'), out_specs=P('x', 'y'))(input)
 
 
 def gradient_kernel(kvec, direction, order=1):
@@ -60,11 +86,7 @@ def laplace_kernel(kvec):
     Complex kernel
   """
     kk = sum(ki**2 for ki in kvec)
-    mask = (kk == 0).nonzero()
-    kk[mask] = 1
-    wts = 1. / kk
-    imask = (~(kk == 0)).astype(int)
-    wts *= imask
+    wts = jnp.where(kk == 0, 1., 1. / kk)
     return wts