split types, improve typing and fix publication download

cosmicweb_music/cosmICweb.py

@@ -2,7 +2,8 @@ import os
 import sys
 import tempfile
 import subprocess
-from typing import NamedTuple, Any, List, Dict
+from typing import Any
+from .data_types import Ellipsoid, Args, DownloadConfig
 
 import click
 import requests
@@ -25,34 +26,6 @@ EDITOR = os.environ.get("EDITOR", "vim")
 EDITOR_IS_VIM = EDITOR in {"vim", "nvim"}
 
 
-# Types
-class Ellipsoid(NamedTuple):
-    center: int
-    shape: int
-    traceback_radius: int
-    radius_definition: int
-
-
-class DownloadConfig(NamedTuple):
-    simulation_name: str
-    project_name: str
-    halo_names: List[str]
-    halo_ids: List[int]
-    halo_urls: List[str]
-    traceback_radius: float
-    api_token: str
-    MUSIC: Dict[str, str]
-    settings: Dict[Any, Any]
-    accessed_at: datetime
-
-
-class Args(NamedTuple):
-    url: str
-    output_path: str
-    common_directory: str
-    attempts: int
-
-
 def query_yes_no(question, default="yes"):
     """Ask a yes/no question via raw_input() and return their answer.
 
@@ -161,6 +134,7 @@ def fetch_publication(cosmicweb_url, publication_name, traceback_radius):
     content = r.json()
     sim = content["simulation"]
     halo_names = [h["name"] for h in content["halos"]]
+    halo_ids = [h["id"] for h in content["halos"]]
     halo_urls = [
         "{url}/simulation/{sid}/halo/{hid}".format(
             url=sim["api_url"], sid=sim["api_id"], hid=h["id"]
@@ -171,12 +145,12 @@ def fetch_publication(cosmicweb_url, publication_name, traceback_radius):
         simulation_name=sim["name"],
         project_name=sim["project_name"],
         halo_names=halo_names,
-        halo_ids=content["halos"],
+        halo_ids=halo_ids,
         halo_urls=halo_urls,
         traceback_radius=traceback_radius,
         api_token=sim["api_token"],
         MUSIC=sim["ics"],
-        settings={},
+        settings=None,
         accessed_at=datetime.now(),
     )
 
@@ -210,27 +184,31 @@ def music_config_to_template(config: DownloadConfig):
     settings = config.settings
     # TODO: apply output configuration
     config = (
-        f"# Zoom Initial Conditions for halo {config.halo_ids[0]} in simulation {config.simulation_name} ({config.project_name} project)\n"
-        f"# Details on this halo can be found on https://cosmicweb.eu/simulation/{config.simulation_name}/halo/{config.halo_ids[0]}\n"
-        f"# This file has been generated by CosmICweb @{datetime.now().isoformat()}\n\n\n"
         "[setup]\n" + music_config["setup"] + "\n\n<ELLIPSOID_TEMPLATE>\n\n"
         "[cosmology]\n" + music_config["cosmology"] + "\n\n"
         "[random]\n" + music_config["random"] + "\n\n"
         "[poisson]\n" + music_config["poisson"]
     )
-    config = apply_config_parameter(
-        config,
-        {
-            "levelmin": settings["resolution"]["low"],
-            "levelmin_TF": settings["resolution"]["low"],
-            "levelmax": settings["resolution"]["high"],
-            "zstart": settings["startRedshift"],
-        },
-    )
+    if settings:
+        config = apply_config_parameter(
+            config,
+            {
+                "levelmin": settings["resolution"]["low"],
+                "levelmin_TF": settings["resolution"]["low"],
+                "levelmax": settings["resolution"]["high"],
+                "zstart": settings["startRedshift"],
+            },
+        )
     return config
 
 
-def compose_template(template, ellipsoid):
+def compose_template(
+    template: str,
+    ellipsoid: Ellipsoid,
+    config: DownloadConfig,
+    halo_name: str,
+    halo_id: int,
+):
     # TODO: add ellipsoid header (rtb, halo_name, etc)
     shape_0 = ", ".join(str(e) for e in ellipsoid.shape[0])
     shape_1 = ", ".join(str(e) for e in ellipsoid.shape[1])
@@ -247,10 +225,16 @@ def compose_template(template, ellipsoid):
         f"region_ellipsoid_matrix[2] = {shape_2}\n"
         f"region_ellipsoid_center    = {center}\n"
     )
-    return template.replace("<ELLIPSOID_TEMPLATE>", ellipsoid_lines)
+    template = template.replace("<ELLIPSOID_TEMPLATE>", ellipsoid_lines)
+    config_header = (
+        f"# Zoom Initial Conditions for halo {halo_id} ({halo_name}) in simulation {config.simulation_name} ({config.project_name} project)\n"
+        f"# Details on this halo can be found on https://cosmicweb.eu/simulation/{config.simulation_name}/halo/{halo_id}\n"
+        f"# This file has been generated by CosmICweb @{datetime.now().isoformat()}\n\n\n"
+    )
+    return config_header + template
 
 
-def write_music_file(output_file, music_config):
+def write_music_file(output_file, music_config) -> None:
     dirname = os.path.dirname(output_file)
     if not os.path.exists(dirname):
         logging.debug("Creating directory {}".format(dirname))
@@ -263,7 +247,7 @@ def call_music():
     pass
 
 
-def process_config(config, args: Args):
+def process_config(config: DownloadConfig, args: Args):
     ellipsoids = []
     for halo_name, url in zip(config.halo_names, config.halo_urls):
         logging.info("Fetching ellipsoids from halo " + halo_name)
@@ -289,14 +273,18 @@ def process_config(config, args: Args):
         music_template = edit_template(music_template)
         logging.debug("Finished editing MUSIC template")
     # Store template to file
-    for halo_name, ellipsoid in zip(config.halo_names, ellipsoids):
+    for halo_name, halo_id, ellipsoid in zip(
+        config.halo_names, config.halo_ids, ellipsoids
+    ):
         if ellipsoid is None:
             logging.warning(
                 "Ellipsoid for halo {} not available, skipping".format(halo_name)
             )
             continue
         logging.info("Composing MUSIC configuration file for halo {}".format(halo_name))
-        music_config = compose_template(music_template, ellipsoid)
+        music_config = compose_template(
+            music_template, ellipsoid, config, halo_name, halo_id
+        )
         if args.common_directory and len(ellipsoids) > 1:
             output_file = os.path.join(args.output_path, str(halo_name), "ics.cfg")
         else:

cosmicweb_music/data_types.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+from typing import NamedTuple, Any, List, Dict, TypedDict
+
+
+class Ellipsoid(NamedTuple):
+    center: List[float]
+    shape: List[List[float]]
+    traceback_radius: float
+    radius_definition: str
+
+
+class Resolution(TypedDict):
+    low: int
+    high: int
+
+
+class Configuration(TypedDict):
+    outputType: str
+    resolution: Resolution
+    outputOptions: List[Any]
+    startRedshift: int
+    outputFilename: str
+    separateFolders: bool
+    tracebackRadius: int | float | str
+
+
+class DownloadConfig(NamedTuple):
+    simulation_name: str
+    project_name: str
+    halo_names: List[str]
+    halo_ids: List[int]
+    halo_urls: List[str]
+    traceback_radius: float
+    api_token: str
+    MUSIC: Dict[str, str]
+    settings: Configuration | None
+    accessed_at: datetime
+
+
+class Args(NamedTuple):
+    url: str
+    output_path: str
+    common_directory: str
+    attempts: int