clm-ml-jax 0.1.0

JAX translation of CLM-ML (Community Land Model with Multi-Layer Canopy)

CLM-ML-JAX: Community Land Model with Multi-Layer Canopy in JAX

A complete JAX/Python translation of the Community Land Model (CLM) Multi-Layer Canopy physics, enabling GPU acceleration

Reference CLM-ML in Fortran: CLM-ML_v2.CHATS

Features

• Full Fortran→Python/JAX translation maintaining scientific accuracy
• Tower-site offline driver for point-scale simulations with prescribed meteorology
• GPU-accelerated physics through JAX JIT compilation
• Vectorization with vmap for ensemble and sensitivity analyses
• Namelist-driven configuration (mirrors Fortran interface)
• Modular architecture

---

Installation

# 1. Clone and navigate to repository
git clone https://github.com/AyaLahlou/clm-ml-jax.git
cd clm-ml-jax

# 2. Install package and dependencies
pip install -e .

# 3. Verify installation
clm-ml-offline --help

This installs the clm-ml-offline command and all physics modules (clm_share, clm_src_*, multilayer_canopy, etc.).

---

Run a Simulation

Execute a 1-day tower-site simulation with the CHATS7 site:

# Or with explicit namelist argument  
clm-ml-offline src/offline_executable/nl.CHATS7.1day

# Or run directly with Python
python -m src.offline_executable.main input_files/nl.CHATS7.1day

Output files are written to the output_files/ directory.

---

Validation Tests

We compare every JAX routine against golden reference values of the Fortran build:

/path/to/conda/envs/clm-ml-jax/bin/python -m pytest tests/fortran_validation/ -q

Golden JSON files live in tests/fortran_validation/golden_IO/.

---

Development

Adding a New Physics Module

1. Create multilayer_canopy/MyPhysicsMod.py mirroring Fortran module structure
2. Include Fortran source reference in docstring (if applicable).
3. Use module-level globals from MLclm_varctl.py for configuration
4. Add tests in tests/multilayer_canopy/test_my_physics.py

Adding a New Tower Site

Edit offline_driver/TowerDataMod.py:

ntower = 16  # Increment from 15

# Extend all arrays by one element
tower_id[16] = 'MYNEWSITE'
tower_lon[16] = -120.5
tower_lat[16] = 38.2
tower_elev[16] = 500.0
# ... (add other fields)

Then create a namelist file: src/offline_executable/nl.MYNEWSITE.date

Configuration & Switches

Global physics switches are in multilayer_canopy/MLclm_varctl.py—change directly, no config object needed:

gs_type = 0              # Stomatal model: 0=Medlyn, 1=Ball-Berry, 2=WUE
flux_profile_type = 1    # Flux-profile: -1=dataset, 0=well-mixed, 1=implicit
runge_kutta_type = 41    # Time integration: 10=Euler, 21=2nd-order, 41=4th-order
dtime_ml = 60.0          # Sub-step interval (s); must divide CLM timestep

---

Contributing

Contributions are welcome! Here's how you can help:

1. Report Bugs: Open an issue with reproduction steps
2. Suggest Features: Open an issue with use case description
3. Submit PRs: Fork, create feature branch, submit PR
4. Improve Docs: Documentation improvements are always welcome

Code of Conduct

• Be respectful and inclusive
• Provide constructive feedback
• Focus on improving the project

---

Citation

@software{clm_ml_jax,
  title={CLM-ML-JAX: Community Land Model with Multi-Layer Canopy in JAX},
  author={Lahlou, Aya},
  year={2024},
  url={https://github.com/AyaLahlou/clm-ml-jax},
  note={JAX translation of CTSM CLM-ML}
}

---

License

This project is licensed under the BSD-3-Clause License. See LICENSE for details.

---

⭐ If this helps your research, please star the repository!

🐛 Found a bug? Please open an issue