musica 0.12.2__cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl

musica/tools/prepare_build_environment_linux.sh

@@ -0,0 +1,32 @@
+#!/bin/bash
+set -e
+set -x
+
+target_arch="$(uname -m)"
+echo "Detected target_arch: $target_arch"
+
+dnf -y update
+
+# For 64 bit systems can enable our fortran components, but we require netcdf
+if [[ "$target_arch" == "x86_64" || "$target_arch" == "aarch64" ]]; then
+  dnf install -y epel-release
+  dnf install -y netcdf-devel netcdf-fortran-devel
+fi
+
+dnf install -y tree wget zip lapack-devel
+
+# 64 bit intel and amd systems support building cuda
+if [ "$target_arch" = "x86_64" ]; then
+  # Install CUDA 12.8 for x86_64 on AlmaLinux 8 (manylinux_2_28) - supports GCC 14
+  dnf config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-rhel8.repo
+  dnf install --setopt=obsoletes=0 -y \
+      cuda-nvcc-12-8 \
+      cuda-cudart-devel-12-8 \
+      libcurand-devel-12-8 \
+      libcublas-devel-12-8 
+  ln -sf cuda-12.8 /usr/local/cuda
+
+  # Verify CUDA installation
+  echo "=== CUDA Installation Verification ==="
+  /usr/local/cuda/bin/nvcc --version
+fi

musica/tools/prepare_build_environment_macos.sh

@@ -0,0 +1 @@
+brew install netcdf netcdf-fortran lapack

musica/tools/repair_wheel_gpu.sh

@@ -0,0 +1,40 @@
+#!/usr/bin/env bash
+set -e
+
+auditwheel repair --exclude libcublas --exclude libcublasLt --exclude libcudart -w "$2" "$1"
+
+for whl in "$2"/*.whl; do
+  tmpdir=$(mktemp -d)
+  unzip -q "$whl" -d "$tmpdir"
+  tree "$tmpdir"
+
+  so_files=("$tmpdir"/musica/_musica_gpu*.so)
+
+  if [[ -f "${so_files[0]}" ]]; then
+      so_path="${so_files[0]}"
+      ls $so_path
+      echo "Before patchelf:"
+      readelf -d $so_path
+      # Use patchelf to fix the rpath and library dependencies
+      patchelf --remove-rpath $so_path
+      patchelf --set-rpath "\$ORIGIN:\$ORIGIN/../nvidia/cublas/lib:\$ORIGIN/../nvidia/cuda_runtime/lib" --force-rpath $so_path
+      # these may need to be periodically updated
+      patchelf --replace-needed libcudart-b5a066d7.so.12.2.140 libcudart.so.12 $so_path
+      patchelf --replace-needed libcublas-e779a79d.so.12.2.5.6 libcublas.so.12 $so_path
+      patchelf --replace-needed libcublasLt-fbfbc8a1.so.12.2.5.6 libcublasLt.so.12 $so_path
+      # Remove bundled CUDA libraries
+      rm -f "$tmpdir"/musica.libs/libcudart-*.so*
+      rm -f "$tmpdir"/musica.libs/libcublas-*.so*
+      rm -f "$tmpdir"/musica.libs/libcublasLt-*.so*
+      echo "After patchelf:"
+      readelf -d $so_path
+  else
+    echo "No GPU .so file found, skipping patchelf steps"
+  fi
+
+  # Repack the wheel with correct structure
+  (cd "$tmpdir" && zip -qr "${whl%.whl}.patched.whl" .)
+  rm -rf "$tmpdir"
+  # Replace the original wheel with the patched one
+  mv "${whl%.whl}.patched.whl" "$whl"
+done

musica/tuvx.cpp

@@ -0,0 +1,93 @@
+#include "binding_common.hpp"
+
+#include <musica/tuvx/tuvx.hpp>
+
+#include <pybind11/numpy.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+namespace py = pybind11;
+
+void bind_tuvx(py::module_& tuvx)
+{
+  tuvx.def("_get_tuvx_version", []() { return musica::TUVX::GetVersion(); }, "Get the version of the TUV-x instance");
+
+  tuvx.def(
+      "_create_tuvx",
+      [](const char* config_path)
+      {
+        try
+        {
+          auto tuvx_instance = new musica::TUVX();
+          tuvx_instance->CreateFromConfigOnly(config_path);
+          return reinterpret_cast<std::uintptr_t>(tuvx_instance);
+        }
+        catch (const std::exception& e)
+        {
+          throw py::value_error(
+              "Error creating TUV-x instance from config file: " + std::string(config_path) + " - " + e.what());
+        }
+      },
+      "Create a TUV-x instance from a JSON configuration file");
+
+  tuvx.def(
+      "_delete_tuvx",
+      [](std::uintptr_t tuvx_ptr)
+      {
+        musica::TUVX* tuvx_instance = reinterpret_cast<musica::TUVX*>(tuvx_ptr);
+        delete tuvx_instance;
+      },
+      "Delete a TUV-x instance");
+
+  tuvx.def(
+      "_run_tuvx",
+      [](std::uintptr_t tuvx_ptr)
+      {
+        musica::TUVX* tuvx_instance = reinterpret_cast<musica::TUVX*>(tuvx_ptr);
+
+        // Get dimensions
+        int n_photolysis = tuvx_instance->GetPhotolysisRateCount();
+
+        int n_heating = tuvx_instance->GetHeatingRateCount();
+
+        int n_layers = tuvx_instance->GetNumberOfLayers();
+
+        int n_sza_steps = tuvx_instance->GetNumberOfSzaSteps();
+
+        // Allocate output arrays (3D: sza_step, layer, reaction/heating_type)
+        std::vector<double> photolysis_rates(n_sza_steps * n_layers * n_photolysis);
+        std::vector<double> heating_rates(n_sza_steps * n_layers * n_heating);
+
+        // Run TUV-x (everything comes from the JSON config)
+        tuvx_instance->RunFromConfig(photolysis_rates.data(), heating_rates.data());
+
+        // Return as numpy arrays with shape (n_sza_steps, n_layers, n_reactions/n_heating)
+        py::array_t<double> py_photolysis =
+            py::array_t<double>({ n_sza_steps, n_layers, n_photolysis }, photolysis_rates.data());
+        py::array_t<double> py_heating = py::array_t<double>({ n_sza_steps, n_layers, n_heating }, heating_rates.data());
+
+        return py::make_tuple(py_photolysis, py_heating);
+      },
+      "Run TUV-x (all parameters come from JSON config)",
+      py::arg("tuvx_instance"));
+
+  tuvx.def(
+      "_get_photolysis_rate_names",
+      [](std::uintptr_t tuvx_ptr)
+      {
+        musica::TUVX* tuvx_instance = reinterpret_cast<musica::TUVX*>(tuvx_ptr);
+
+        return tuvx_instance->GetPhotolysisRateNames();
+      },
+      "Get photolysis rate names");
+
+  tuvx.def(
+      "_get_heating_rate_names",
+      [](std::uintptr_t tuvx_ptr)
+      {
+        musica::TUVX* tuvx_instance = reinterpret_cast<musica::TUVX*>(tuvx_ptr);
+
+        return tuvx_instance->GetHeatingRateNames();
+      },
+      "Get heating rate names");
+}

musica/tuvx.py

@@ -0,0 +1,199 @@
+"""
+TUV-x photolysis calculator Python interface.
+
+This module provides a simplified Python interface to the TUV-x photolysis calculator.
+It allows users to create a TUV-x instance from a JSON configuration file and
+calculate photolysis rates and heating rates.
+
+Note: TUV-x is only available on macOS and Linux platforms.
+"""
+
+import os
+import json
+import tempfile
+from typing import Dict, Tuple, List
+import numpy as np
+from . import backend
+
+_backend = backend.get_backend()
+
+version = _backend._tuvx._get_tuvx_version() if backend.tuvx_available() else None
+
+
+class TUVX:
+    """
+    A Python interface to the TUV-x photolysis calculator.
+
+    This class provides a simplified interface that only requires a JSON configuration
+    file to set up and run TUV-x calculations. All parameters (solar zenith angle,
+    earth-sun distance, atmospheric profiles, etc.) are specified in the JSON config.
+    """
+
+    def __init__(self, config_path: str):
+        """
+        Initialize a TUV-x instance from a configuration file.
+
+        Args:
+            config_path: Path to the JSON configuration file
+
+        Raises:
+            FileNotFoundError: If the configuration file doesn't exist
+            ValueError: If TUV-x initialization fails or TUVX is not available
+        """
+        if not backend.tuvx_available():
+            raise ValueError("TUV-x backend is not available on windows.")
+
+        if not os.path.exists(config_path):
+            raise FileNotFoundError(
+                f"Configuration file not found: {config_path}")
+
+        self._tuvx_instance = _backend._tuvx._create_tuvx(config_path)
+        self._config_path = config_path
+
+        # Cache the names for efficiency
+        self._photolysis_names = None
+        self._heating_names = None
+
+    def __del__(self):
+        """Clean up the TUV-x instance."""
+        if hasattr(self, '_tuvx_instance') and self._tuvx_instance is not None:
+            _backend._tuvx._delete_tuvx(self._tuvx_instance)
+
+    @property
+    def photolysis_rate_names(self) -> List[str]:
+        """
+        Get the names of photolysis rates.
+
+        Returns:
+            List of photolysis rate names
+        """
+        if self._photolysis_names is None:
+            self._photolysis_names = _backend._tuvx._get_photolysis_rate_names(
+                self._tuvx_instance)
+        return self._photolysis_names
+
+    @property
+    def heating_rate_names(self) -> List[str]:
+        """
+        Get the names of heating rates.
+
+        Returns:
+            List of heating rate names
+        """
+        if self._heating_names is None:
+            self._heating_names = _backend._tuvx._get_heating_rate_names(
+                self._tuvx_instance)
+        return self._heating_names
+
+    def run(self) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        Run the TUV-x photolysis calculator.
+
+        All parameters (solar zenith angle, Earth-Sun distance, atmospheric profiles,
+        etc.) are read from the JSON configuration file.
+
+        Returns:
+            Tuple of (photolysis_rate_constants, heating_rates) as numpy arrays
+            - photolysis_rate_constants: Shape (n_layers, n_reactions) [s^-1]
+            - heating_rates: Shape (n_layers, n_heating_rates) [K s^-1]
+        """
+        photolysis_rates, heating_rates = _backend._tuvx._run_tuvx(
+            self._tuvx_instance)
+
+        return photolysis_rates, heating_rates
+
+    def get_photolysis_rate_constant(
+        self,
+        reaction_name: str,
+        photolysis_rates: np.ndarray
+    ) -> np.ndarray:
+        """
+        Extract photolysis rate constants for a specific reaction.
+
+        Args:
+            reaction_name: Name of the photolysis reaction
+            photolysis_rates: Output from run() method
+
+        Returns:
+            1D array of photolysis rate constants for all layers [s^-1]
+
+        Raises:
+            KeyError: If reaction_name is not found
+        """
+        names = self.photolysis_rate_names
+        if reaction_name not in names:
+            raise KeyError(
+                f"Reaction '{reaction_name}' not found. "
+                f"Available reactions: {names}"
+            )
+
+        reaction_index = names.index(reaction_name)
+        return photolysis_rates[:, reaction_index]
+
+    def get_heating_rate(
+        self,
+        rate_name: str,
+        heating_rates: np.ndarray
+    ) -> np.ndarray:
+        """
+        Extract heating rates for a specific rate type.
+
+        Args:
+            rate_name: Name of the heating rate
+            heating_rates: Output from run() method
+
+        Returns:
+            1D array of heating rates for all layers [K s^-1]
+
+        Raises:
+            KeyError: If rate_name is not found
+        """
+        names = self.heating_rate_names
+        if rate_name not in names:
+            raise KeyError(
+                f"Heating rate '{rate_name}' not found. "
+                f"Available rates: {names}"
+            )
+
+        rate_index = names.index(rate_name)
+        return heating_rates[:, rate_index]
+
+    @staticmethod
+    def create_config_from_dict(config_dict: Dict) -> 'TUVX':
+        """
+        Create a TUVX instance from a configuration dictionary.
+
+        Args:
+            config_dict: Configuration dictionary
+
+        Returns:
+            TUVX instance initialized with the configuration
+
+        Raises:
+            ValueError: If TUV-x backend is not available
+            FileNotFoundError: If required data files are not found
+        """
+        with tempfile.NamedTemporaryFile(
+                mode='w', suffix='.json', delete=True) as temp_file:
+            json.dump(config_dict, temp_file, indent=2)
+            temp_file.flush()  # Ensure all data is written to disk
+            return TUVX(temp_file.name)
+
+    @staticmethod
+    def create_config_from_json_string(json_string: str) -> 'TUVX':
+        """
+        Create a TUVX instance from a JSON configuration string.
+
+        Args:
+            json_string: JSON configuration as string
+
+        Returns:
+            TUVX instance initialized with the configuration
+
+        Raises:
+            json.JSONDecodeError: If json_string is not valid JSON
+            ValueError: If TUV-x backend is not available
+            FileNotFoundError: If required data files are not found
+        """
+        config_dict = json.loads(json_string)
+        return TUVX.create_config_from_dict(config_dict)