virgodc 1.0.2
Tools for working with Virgo consortium n-body simulations.
Requires Python: >=3.10

# Virgo python module

## Introduction

This module provides facilities for reading various formats used to store
snapshots, group catalogues and merger trees in Virgo Consortium simulations,
including older binary formats which can otherwise be difficult to deal with.

It also provides a collection of MPI parallel algorithms which can be useful
for dealing with particle data and a few related utility functions.

## Installation

This module requires mpi4py and h5py. On a HPC system these may need to be
built from source to ensure they're linked to the right MPI implementation.

To install mpi4py, ensure that the right mpicc is in your $PATH (e.g. by
loading environment modules) and run
```
python -m pip install mpi4py
```

To install h5py, if the right mpicc is in your $PATH and HDF5 is installed at
$HDF5_HOME:
```
export CC="mpicc"
export HDF5_MPI="ON"
export HDF5_DIR=${HDF5_HOME}
pip install --no-binary=h5py h5py
```

Running the tests requires pytest-mpi:
```
pip install pytest-mpi
```

The module can then be installed using pip:
```
pip install virgodc
```

## Layout of the module

  * virgo.formats: contains classes for reading various simulation data formats
  * virgo.sims: contains wrappers for the read routines with appropriate default parameters for particular simulations
  * virgo.util: contains utility functions which may be useful when working with simulation data, including
    * the BinaryFile class, which allows for reading binary files using a HDF5-like interface
    * an efficient method for finding matching values in arrays of particle IDs
    * a vectorized python implementation of the peano_hilbert_key function from Gadget
  * virgo.mpi: utilities for working with simulation data using mpi4py, including a reasonably efficient MPI parallel sort and functions for parallel I/O on multi-file simulation output

## Supported data formats

### SWIFT snapshots

For a more complete solution for reading SWIFT snapshots, see
[swiftsimio](https://swiftsimio.readthedocs.io/en/latest/).

This module provides a simple class for reading SWIFT snapshots which is
essentially a h5py.File that returns arrays with units attached when datasets
are read. A snapshot can be opened as follows:
```
from virgo.formats.swift import SwiftSnapshot
snap = SwiftSnapshot("snapshot_0010.hdf5")
```
To read a dataset we can index the snapshot object as with a h5py.File:
```
>>> pos = snap["PartType0/Coordinates"][...]
>>> print(pos)
[[ 34.47785792 565.96551792 409.66667792]
 [ 32.55394792 563.85734792 409.80881792]
 [ 32.65551792 562.94646792 408.11432792]
 ...
 [342.2256186  812.0977686  468.7078186 ]
 [342.0442786  812.3124186  467.4161186 ]
 [342.0407086  811.3151886  467.6265686 ]] a*snap_length)
```

#### Units

When reading a snapshot the module defines several custom unyt units:
  * `snap_length`: the length unit used in the snapshot
  * `snap_mass`: the mass unit used in the snapshot
  * `snap_time`: the time unit used in the snapshot
  * `snap_temperature`: the temperature unit used in the snapshot
  * `a` : the expansion factor of this snapshot
  * `h` : the Hubble parameter used in the simulation

Reading a dataset returns a unyt array using these base units. A comoving
position, for example, would have units `a*snap_length` and a velocity would
have units `snap_length/snap_time`.

It also defines several symbols based on SWIFT's physical constants:
  * `swift_mpc` : SWIFT's definition of a megaparsec
  * `swift_msun` : SWIFT's definition of a solar mass
  * `newton_G` : the gravitational constant used in the simulation

To return positions in comoving Mpc regardless of the simulation unit system,
for example, we could do:
```
pos = snap["PartType0/Coordinates"][...].to("a*swift_mpc")
```
This avoids any slight changes to the values due to SWIFT and unyt having
different definitions of a parsec or the mass of the sun.

#### Cosmology

The SwiftSnapshot class has a cosmology field which returns an astropy
cosmology object based on the parameters in the simulation snapshot.

### SOAP halo catalogues

Halo catalogues generated using the SOAP tool can be read in a similar way:
```
from virgo.formats.soap import SOAPCatalogue
cat = SOAPCatalogue("halo_properties_0077.hdf5")
total_mass = cat["SO/200_crit/TotalMass"][...]
```
This returns a unyt array using unit scheme described above. The cosmology
can be accessed in the same way too.

### Gadget snapshots

The function virgo.formats.gadget_snapshot.open() can be used to read Gadget snapshots
stored in either HDF5 or type 1 binary format. When opening a snapshot file it returns
either a h5py.File object (if the file is in HDF5 format) or a GadgetSnapshotFile
object (if the file is in binary format). In the case of binary files, the precision 
and endian-ness of the file are determined automatically.

Quantities in binary snapshots are accessed using HDF5 style names:
```
import virgo.formats.gadget_snapshot as gs
snap = gs.open("snap_C02_400_1023.0")
boxsize = snap["Header"].attrs["BoxSize"]
pos     = snap["PartType1/Coordinates"][...]
vel     = snap["PartType1/Velocities"][...]
```

### Gadget Friends-of-Friends and Subfind output

The following modules contains classes to read subfind and friends of friends output from several versions of Gadget:

  * virgo.formats.subfind_lgadget2 - L-Gadget2, e.g. the Millennium simulation
  * virgo.formats.subfind_lgadget3 - L-Gadget3, e.g. MXXL
  * virgo.formats.subfind_pgadget3 - Millennium2, Aquarius
  * virgo.formats.subfind_gadget4  - The version of Gadget-4 used in COCO (likely not compatible with the current Gadget-4)

These each provide the following classes:

  * SubTabFile - to read subfind group catalogues
  * SubIDsFile - to read the IDs of particles in subfind groups
  * GroupTabFile - to read friends of friends catalogues
  * GroupIDsFile - to read the IDs of particles in friends of friends groups

In some cases extra parameters are required before files can be read:

  * id_bytes - number of bytes used to store a particle ID (4 or 8)
  * float_bytes - number of bytes used to store float quantities in group catalogues (4 or 8)
  * Flags corresponding to Gadget pre-processor macros which affect the output format. These
    should be set to True or False depending on whether the corresponding macro was set.
    * SO_VEL_DISPERSIONS
    * SO_BAR_INFO

See the docstrings associated with each class to determine which parameters are required for which formats.
In most cases calling the sanity_check() method on the object will raise an exception if any parameters
were set incorrectly.

## MPI parallel algorithms

Working with halo finder output often requires matching particle IDs between
the halo finder output files and simulation snapshots.

The module virgo.mpi.parallel_sort contains functions for sorting, matching
and finding unique values in distributed arrays. A distributed array is simply
a numpy array which exists on each MPI rank and is treated as forming part of
a single, large array. Elements in distributed arrays have a notional 'global'
index which, for an array with N elements over all ranks, runs from zero for
the first element on rank zero to N-1 for the last element on the last MPI
rank.

All of these functions are collective and must be executed an all MPI ranks in
the specified communicator `comm`, or `MPI_COMM_WORLD` if the `comm` parameter
is not supplied.

### Repartitioning a distributed array

```
virgo.mpi.parallel_sort.repartition(arr, ndesired, comm=None)
```
This function returns a copy of the input distributed array `arr` with the
number of elements per MPI rank specified in ndesired, which must be an integer
array with one element per MPI rank.

This can be used to improve memory load balancing if the input array is
unevenly distrubuted. `comm` specifies the MPI communicator to use.
MPI_COMM_WORLD is used if comm is not set.

If `arr` is multidimensional then the repartitioning is done along the first
axis.

### Fetching specified elements of a distributed array

```
virgo.mpi.parallel_sort.fetch_elements(arr, index, result=None, comm=None)
```
This function returns a new distributed array containing the elements of `arr`
with global indexes specified by `index`. On one MPI rank the result would just
be `arr[index]`. The output can be placed in an existing array using the
`result` parameter.

This function can be used to apply the sorting index returned by
parallel_sort().

If `arr` is multidimensional then the index is taken to refer to the first
axis. I.e. if running with only one MPI rank the function would return
`arr[index,...]`.

### Sorting

```
virgo.mpi.parallel_sort.parallel_sort(arr, comm=None, return_index=False, verbose=False)
```
This function sorts the supplied array `arr` in place. If `return_index` is
true then it also returns a new distributed array with the global indexes
in the input array of the returned values.

When `return_index` is used this is essentially an MPI parallel version of
np.argsort but with the side effect of sorting the array in place. The index
can be supplied to fetch_elements() to put other arrays into the same order
as the sorted array.

Only one dimensional arrays can be sorted.

This is an attempt to implement the sorting algorithm described in
https://math.mit.edu/~edelman/publications/scalable_parallel.pdf in python,
although there are some differences. For example, the final merging of sorted
array sections is done using a full sort due to the lack of an optimized merge
function in numpy.

### Finding matching values between two arrays

```
virgo.mpi.parallel_sort.parallel_match(arr1, arr2, arr2_sorted=False, comm=None)
```
For each value in the input distributed array `arr1` this function returns the
global index of an element with the same value in distributed array `arr2`, or
-1 where no match is found.

If `arr2_sorted` is True then the input `arr2` is assumed to be sorted, which
saves some time. Incorrect results will be generated if `arr2_sorted` is true
but `arr2` is not really sorted.

Both input arrays must be one dimensional.

### Finding unique values

```
virgo.mpi.parallel_sort.parallel_unique(arr, comm=None, arr_sorted=False,
    return_counts=False, repartition_output=False)
```

This function returns a new distributed array which contains the unique values
from the input distributed array arr. If `arr_sorted` is True the input is
assumed to already be sorted, which saves some time.

If `return_counts` is true then it also returns a distributed array with the
number of instances of each unique value which were found.

If `repartition_output` is true then the output arrays are repartitioned to
have approximately equal numbers of elements on each MPI rank.

The input array must be one dimensional.

### Parallel bin count

```
virgo.mpi.parallel_sort.parallel_bincount(x, weights=None, minlength=None,
                                          result=None, comm=None)
```

This is roughly equivalent to np.bincount but for distributed arrays. Given
an input 1D array x with integer values 0 to N it counts the number of
instances of each integer and returns the result as a new distributed array.

If `weights` is specified then it must be an array with the same number of
elements as x. The result will then be the sum of the weights associated
with each integer value.

`minlength` specifies the minimum size of the output array and `result`
allows the function to write its output to an existing array.

### MPI Alltoallv function

```
virgo.mpi.parallel_sort.my_alltoallv(sendbuf, send_count, send_offset,
    recvbuf, recv_count, recv_offset,
    comm=None)
```

Most of the operations in this module use all-to-all communication patterns.
This function provides an all-to-all implementation that avoids problems
with large (>2GB) communications and can handle any numpy type that the mpi4py
`mpi4py.util.dtlib.from_numpy_dtype()` function can translate into an MPI type.

  * `sendbuf` - numpy array with the data to send
  * `send_count` - number of elements to go to each MPI rank
  * `send_offset` - offset of the first element to send to each rank
  * `recvbuf` - numpy array to receive data into
  * `recv_count` - number of elements to go to receive from MPI rank
  * `recv_offset` - offset of the first element to receive from each rank
  * `comm` - specifes the communicator to use (MPI_COMM_WORLD if not set)

### Tests

The parallel_sort module includes several tests, which can be run with
pytest-mpi:

```
cd VirgoDC/python
mpirun -np 8 python3 -m pytest --with-mpi
```

A larger parallel sort test can be run with
```
mpirun -np 16 python3 -m mpi4py -c "import virgo.mpi.test_parallel_sort as tps ; tps.run_large_parallel_sort(N)"
```
where N is the number of elements per rank to sort. This sorts an array of
random integers and checks that the result is in order and contains the same
number of instances of each value as the input.

## MPI I/O Functions

The module virgo.mpi.parallel_hdf5 contains functions for reading and writing
distributed arrays stored in sets of HDF5 files, using MPI collective I/O
where possible. These can be useful for reading simulation snapshots and halo
finder output.

### Collective HDF5 Read

```
virgo.mpi.parallel_hdf5.collective_read(dataset, comm)
```
This function carries out a collective read of the dataset `dataset`, which
should be a h5py.Dataset object in a file opened with the 'mpio' driver. It
returns a new distributed array with the data.

Multidimensional arrays are partitioned between MPI ranks along the first
axis.

Reads are chunked if necessary to avoid problems with the underlying MPI
library failing to handle reads of >2GB.

### Collective HDF5 Write

```
virgo.mpi.parallel_hdf5.collective_write(group, name, data, comm, create_dataset=True)
```
This function writes the distributed array `data` to the h5py.Group specified
 by the `group` parameter with name `name`. If create_dataset is True then a
new dataset will be created. Otherwise, it is assumed that a dataset with
suitable type and dimensions already exists.

Multidimensional arrays are assumed to be distributed between MPI ranks along
the first axis. Writes are chunked if necessary to avoid problems with the
underlying MPI library failing to handle writes of >2GB.

Returns the new (or modified) h5py.Dataset object.

### Multi-file Parallel I/O

#### The MultiFile class

```
virgo.mpi.parallel_hdf5.MultiFile.__init__(self, filenames, file_nr_attr=None,
    file_nr_dataset=None, file_idx=None, comm=None)
```

Simulation codes can often split their output over a variable number of files.
There may be a single large output file, many small files, or something in
between. This class is intended to solve the general problem of carrying out
parallel reads of distributed arrays from N files on M MPI ranks for arbitrary
values of N and M.

The approach is as follows:

  * For N >= M (i.e. at least one file per MPI rank) each MPI rank uses serial
    I/O to read a subset of the files
  * For N < M (i.e. more MPI ranks than files) the MPI ranks are split into
    groups and each group does collective I/O on one file

The class takes the following parameters:
  * `filenames` - a format string to generate the names of files in the set,
    or a list of strings with the file names. If a format string the file
    number is substituted in as `filenames % {"file_nr" : file_nr}`
  * `file_nr_attr` - a tuple with (HDF5 object name, attribute name) which
    specifies a HDF5 attribute containing the number of files in the set.
    E.g. in a Gadget snapshot use
    `file_nr_attr=("Header","NumFilesPerSnapshot")`.
  * `file_nr_dataset` - the name of a dataset with the number of files in the
     set
  * `file_idx` - an array with the indexes of the files in the set

Exactly one of `file_nr_attr`, `file_nr_dataset` and `file_idx` must be
specified if `filenames` is not a list of strings.

#### Reading datasets from a file set

```
virgo.mpi.parallel_hdf5.MultiFile.read(self, datasets, group=None,
    return_file_nr=None, read_attributes=False)
```
This method reads one or more distributed arrays from the file set. The arrays
are distributed between MPI ranks along the first axis. The parameters are:
  * `datasets` - a list of the names of the datasets to read, or a string
     specifying a single dataset to read
  * `group` - the name of the HDF5 group to read datasets from
  * `return_file_nr` - if this is true an additional set of arrays is returned
     which contain the index of the file each dataset element was read from.
  * `read_attributes` - if this is true, return an ndarray subclass with a
    .attrs attribute, which is a dict containing all HDF5 attributes of the
    dataset in the file. Attributes are assumed to be the same between files.
  * `unpack` - if True, results that would be returned as dicts are returned
     as lists instead.

This reads the specified datasets and for each one returns a distributed array
with the data. The arrays are returned in one of three ways:

  * If `datasets` is a single string then a single array is returned
  * If `datasets` is a list of strings and unpack=False, returns a dict where
    the keys are dataset names and the values are numpy arrays with the data
  * If `datasets` is a list of strings and unpack=True, returns a list of
    numpy arrays

This can be used to read particles from a snapshot distributed over an
arbitrary number of files, for example.

The unpack option is useful for avoiding repetition of the dataset names. E.g.
```
data = mf.read(("Positions", "Velocities"))
pos = data["Positions"]
vel = data["Velocities"]
```
can be reduced to
```
pos, vel = mf.read(("Positions", "Velocities"), unpack=True)
```

#### Reading the number of dataset elements per file

```
virgo.mpi.parallel_hdf5.MultiFile.get_elements_per_file(self, name, group=None)
```
This returns the number of elements read from each file for the specified
dataset. Note that in collective mode the return value varies between ranks
because each rank returns the number of elements which it will read from the file,
and not the total number of elements in the file.

This should therefore only be used with `MultiFile.write()` to write output
distributed between files in the same way as an input file set.

  * `name` - name of the dataset
  * `group` - name of the group containing the dataset

#### Writing datasets to a file set

```
virgo.mpi.parallel_hdf5.MultiFile.write(self, data, elements_per_file,
    filenames, mode, group=None, attrs=None)
```
This writes the supplied distributed arrays to a set of output files with the
specified number of elements per file. The number of output files is the same
as in the input file set used to initialize the class.

  * `data` - a dict containing the distributed arrays to write out. The dict
    keys are used as output dataset names
  * `elements_per_file` - the number of elements along the first axis to write
    to each output file. In collective mode this is the number of elements to
    write from each rank.
  * `filenames` - a format string to generate the names of files in the set.
    The file number is substituted in as `filenames % {"file_nr" : file_nr}`
  * `mode` - should be 'r+' to write to existing files or 'w' to create new files
  * `group` - the name of the HDF5 group to write the datasets to
  * `attrs` - a dict containing attributes to add to the datasets, of the form
    `attrs[dataset_name] = (attribute_name, attribute_value)`

The get_elements_per_file() method can be used to get the value of
elements_per_file needed to write output partitioned in the same way as some
input file set.

WARNING: it is only safe to use this to write arrays which are distributed
between MPI ranks in the same way as an array which was read using
MultiFile.read(), because it assumes that array elements are already on the
rank which will write the file they should go to. Setting arbitrary values of
elements_per_file will cause incorrect output or a crash.

TODO: make elements_per_file actually reflect the number of elements per file
and implement automatic repartitioning of the input so that arbitrary values
of elements_per_file will work as expected.

## MPI Utility Functions

The module virgo.mpi.util contains several other functions which are helpful
for dealing with simulation and halo finder output.

### Computing particle group membership from Subfind lengths and offsets

```
virgo.mpi.util.group_index_from_length_and_offset(length, offset,
    nr_local_ids, return_rank=False, comm=None)
```

Given distributed arrays with the lengths and offsets of particles in a subfind
output, this computes the group index for each particle. The first group is
assigned index zero.
  * `length` - distributed array with the number of particles in each group
  * `offset` - distributed array with the offset to the first particle in each
    group
  * `nr_local_ids` - size of the particle IDs array on this MPI rank. Used to
    determine the size of the output group membership array
  * `comm` - communicator to use. Will use MPI_COMM_WORLD if not specified.

On one MPI rank this would be equivalent to:
```
grnr = np.ones(nr_local_ids, dtype=int)
for i, (l, o) in enumerate(zip(length, offset)):
  grnr[o:o+l] = i
return grnr
```

This can be used in combination with virgo.mpi.parallel_sort.parallel_match()
to find subfind group membership for particles in a simulation snapshot.

If the parameter return_rank=True then it also returns an array with the 
rank ordering of each particle in its halo, with zero indicating the first
particle in the halo. 

### Allocating zero-sized arrays on ranks with no data

Gadget snapshots typically omit HDF5 datasets which would have zero size (e.g.
if some files in a snapshot happen to have zero star particles). This can be an
issue in parallel programs because MPI ranks which read files with such missing
datasets don't know the type or dimensions of some of the datasets.

```
virgo.mpi.util.replace_none_with_zero_size(arr, comm=None)
```
This takes an input distributed array, `arr`, and on ranks where arr is None
an empty array is returned using type and size information from the other MPI
ranks. The new array will have zero size in the first dimension and the same
size as the other MPI ranks in all other dimensions.

On ranks where the input is not None, the input array is returned.

The array should have the same dtype on all ranks where it is not None.

The intended use of this function is to allow read routines to return None
where datasets do not exist and then this function can be used to retrieve the
missing metadata.

### MPI argument parser

virgo.mpi.util.MPIArgumentParser is a subclass of argparse.ArgumentParser for
use in MPI programs. It parses arguments on rank 0 of the supplied communicator
and broadcasts them to the rest of the communicator. In the event of an error
it prints a message to rank 0's stderr, calls MPI_Finalize, and terminates
all processes.

It takes the communicator to use as its first argument. This should usually be
mpi4py.MPI.COMM_WORLD. Any other parameters are passed to
argparse.ArgumentParser.