Add notebook for galaxy image simulation with Dask

# ---

# jupyter:

#   jupytext:

#     formats: ipynb,py:light

#     text_representation:

#       extension: .py

#       format_name: light

#       format_version: '1.5'

#       jupytext_version: 1.16.2

#   kernelspec:

#     display_name: Python 3 (ipykernel)

#     language: python

#     name: python3

# ---

# # Parallelizing simulations with Dask

#

# Parallelizing the simulation of Galaxies is an example

# of a workload which is very well suited for Dask.

#

# This notebook is an example of how could this workflow be implemented. It does not actually use real parameters or simulator codes, insted, it uses random parameter values, random images and does dummy calculations.

#

# ## Prerequisites

#

# * Packages

#   * dask ecosystem

#     * dask

#     * distributed

#     * dask.jobqueue

#   * tqdm

#   * zarr

# * Running Dask cluster

# +

import math

import random

import time

import numpy as np

import pandas as pd

from pathlib import Path

from tqdm.notebook import tqdm

import dask

import dask.array as da

from dask.distributed import Client

# -

# ## Dask Cluster Connection

#

# Before running the simulations, we connect to a Dask cluster. This allows for distributed and parallel computation, speeding up the overall process. Ensure the Dask cluster is running and accessible from this environment.

#

# In this example the code to connect to the Dask cluster was injected by the dask jupyterlab extension (`dask-labextension`).

#

# The code would change if the cluster is spawned within the notebook. For example:

#

# ```

# from dask_jobqueue import HTCondorCluster

# from dask.distributed import Client

# cluster = HTCondorCluster(cores=1, memory='2GB', disk='10 GB',

#                          job_extra_directives={'getenv': 'True'})

# cluster.scale(10)

# client = Client(cluster)

# ```

# +

from dask.distributed import Client

client = Client("tls://192.168.102.15:32769")

client

# -

# ## Data Simulation Functions

#

# This section defines the core functions used for simulating galaxy images.

# +

def dummy_calculation(duration=1):

    """

    Burns CPU cycles for a specified duration to simulate computational work.

    This function performs a series of mathematical operations in a loop

    until the specified `duration` (in seconds) has passed. It's used

    to mimic a more realistic computation time rather than just pausing

    execution with `time.sleep()`.

    Args:

        duration (float): The minimum duration in seconds for which to run the dummy calculation.

    """

    start_time = time.time()

    # Dummy computation: math operations in a loop

    result = 0.0

    while time.time() - start_time < duration:

        x = random.random()

        y = random.random()

        result += math.sin(x) * math.cos(y) + math.sqrt(x * y)

def sim_galaxies(params):

    """

    Simulates galaxy image data based on input parameters.

    This function takes a DataFrame of galaxy parameters and generates

    simulated image data (represented as NumPy arrays) for each galaxy.

    It includes a `dummy_calculation` to simulate the computational

    effort involved in a real galaxy simulation.

    Args:

        params (pd.DataFrame): A DataFrame where each row represents

                               parameters for a single galaxy simulation.

                               The number of rows determines `Ngal`.

    Returns:

        np.ndarray: A 4D NumPy array of shape (Ngal, 2, 64, 64)

                    containing the simulated galaxy images.

                    Each galaxy has 2 channels and an image size of 64x64.

    """

    Ngal = len(params)

    A = np.zeros((Ngal, 2, 64, 64))

    for i in range(Ngal):

        # 

        dummy_calculation(duration=0.01)

        # Here you would call up a normal Python function to calculate the

        # values.

        sims = np.random.random((2, 64, 64))

        A[i] = sims

    return A

# -

# ## Generating Metadata

#

# We generate some random metadata that will be the input for the galaxy simulation code. This metadata might include properties like flux, amplitude, size, etc.

# +

# Generate some random metadata.

cols = ['flux', 'A', 'size', 'e1', 'e2']

meta = pd.DataFrame(np.random.random((50_000, 5)), columns=cols)

# Memory usage of the metadata.

mem_meta = meta.memory_usage().sum() / 1024**2

print(f'Memory meta (MB): {mem_meta:.2f}')

# -

# To efficiently process the simulations in parallel using Dask, we split the metadata into smaller chunks. Each chunk will correspond to a Dask task, simulating a batch of galaxies

# +

# How many galaxies to calculate in each job. I would prefer having more than

# one since scheduling the jobs have some overhead.

Ncalc = 100

metaL = np.array_split(meta, int(len(meta) / Ncalc) + 1)

# -

# ## Computation definition

#

# The `sim_galaxies` function is wrapped with `dask.delayed` to create a lazy function. This means that when `f_sim` is called, it doesn't execute immediately but instead builds a Dask graph. The results are then converted into Dask arrays (`da.from_delayed`) and concatenated into a single Dask array.

#

# The actual computation takes place later on.

# +

# This could also be put as a decorator on the function, but having it

# separately makes debugging easier.

f_sim = dask.delayed(sim_galaxies)

L = []

for x in tqdm(metaL):

    L.append(da.from_delayed(f_sim(x), shape=(len(x), 2, 64, 64), dtype=np.float64))

comb = da.concatenate(L)

# Call rechunk to prevent errors when exporting to Zarr format

comb = comb.rechunk(chunks='auto')

# -

# You can see the structure (chunking) of the comb Dask Array.

comb

# ## Computing and storing Results

#

# Finally, the simulated galaxy data and its associated metadata are stored. The galaxy simulations, which are Dask arrays, are saved to a Zarr array and the metadata is saved to a Parquet file.

#

# he computations defined earlier are executed when `comb.to_zarr()` is called

# +

d_out = Path('/data/pic/scratch/torradeflot/tmp')

# When asking to store the calculations starts.

comb.to_zarr(d_out / 'galaxy_sim_v2.zarr')

# Also dump the associated metadata.

meta.to_parquet(d_out / 'galaxy_sim_v2.pq')

# -