Metadata-Version: 2.4
Name: clm-ml-jax
Version: 0.1.0
Summary: JAX translation of CLM-ML (Community Land Model with Multi-Layer Canopy)
Author-email: Aya Lahlou <al4385@columbia.edu>
License-Expression: BSD-3-Clause
Project-URL: Homepage, https://github.com/AyaLahlou/clm-ml-jax
Project-URL: Repository, https://github.com/AyaLahlou/clm-ml-jax
Keywords: jax,land-surface-model,canopy,CLM,differentiable,photosynthesis
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering :: Atmospheric Science
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: jax>=0.4.0
Requires-Dist: jaxlib>=0.4.0
Requires-Dist: numpy>=1.20.0
Requires-Dist: matplotlib>=3.7.0
Requires-Dist: f90nml>=1.4.0
Requires-Dist: netCDF4>=1.6.0
Provides-Extra: dev
Requires-Dist: pytest>=7.0.0; extra == "dev"
Requires-Dist: pytest-cov>=4.0.0; extra == "dev"
Requires-Dist: black>=23.0.0; extra == "dev"
Requires-Dist: ruff>=0.1.0; extra == "dev"
Dynamic: license-file

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

[![CI](https://github.com/AyaLahlou/clm-ml-jax/actions/workflows/ci.yml/badge.svg)](https://github.com/AyaLahlou/clm-ml-jax/actions/workflows/ci.yml)
[![codecov](https://codecov.io/gh/AyaLahlou/clm-ml-jax/branch/main/graph/badge.svg)](https://codecov.io/gh/AyaLahlou/clm-ml-jax)
[![Python 3.9+](https://img.shields.io/badge/python-3.9+-blue.svg)](https://www.python.org/downloads/)
[![JAX](https://img.shields.io/badge/JAX-latest-orange.svg)](https://github.com/google/jax)
[![License: BSD-3](https://img.shields.io/badge/License-BSD--3-green.svg)](LICENSE)

> 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](https://github.com/gbonan/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

```bash
# 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:

```bash

# 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:

```bash
/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`:

```python
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:

```python
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


```bibtex
@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](LICENSE) for details.

---

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

🐛 **Found a bug?** Please [open an issue](https://github.com/AyaLahlou/clm-ml-jax/issues)

