Metadata-Version: 2.4
Name: transjax
Version: 0.1.0
Summary: Translate Fortran scientific code to JAX using LLM agents
Author: Aya Lahlou
License: BSD-3-Clause
Project-URL: Homepage, https://github.com/AyaLahlou/TransJAX
Project-URL: Repository, https://github.com/AyaLahlou/TransJAX
Project-URL: Bug Tracker, https://github.com/AyaLahlou/TransJAX/issues
Keywords: fortran,jax,translation,llm,agents,scientific-computing
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: BSD License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering
Classifier: Topic :: Software Development :: Code Generators
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: anthropic>=0.40.0
Requires-Dist: python-dotenv>=1.0.0
Requires-Dist: pyyaml>=6.0
Requires-Dist: pydantic>=2.0.0
Requires-Dist: rich>=13.0.0
Requires-Dist: tenacity>=8.0.0
Requires-Dist: click>=8.0.0
Requires-Dist: networkx>=2.5
Requires-Dist: matplotlib>=3.3.0
Requires-Dist: numpy>=1.20.0
Provides-Extra: testing
Requires-Dist: jax>=0.4.0; extra == "testing"
Requires-Dist: jaxlib>=0.4.0; extra == "testing"
Requires-Dist: pytest>=7.0.0; extra == "testing"
Provides-Extra: dev
Requires-Dist: jax>=0.4.0; extra == "dev"
Requires-Dist: jaxlib>=0.4.0; extra == "dev"
Requires-Dist: pytest>=7.0.0; extra == "dev"
Requires-Dist: black>=23.0.0; extra == "dev"
Requires-Dist: ruff>=0.1.0; extra == "dev"
Requires-Dist: mypy>=1.0.0; extra == "dev"
Requires-Dist: pytest-cov>=4.0.0; extra == "dev"
Requires-Dist: tomli>=2.0.0; python_version < "3.11" and extra == "dev"
Provides-Extra: viz
Requires-Dist: seaborn>=0.11.0; extra == "viz"
Dynamic: license-file

# TransJAX

**Translate Fortran scientific code to JAX automatically using LLM agents.**

TransJAX is a unified Python package that combines static Fortran analysis with
a multi-agent translation pipeline powered by Claude (Anthropic).  It is designed
for the scientific-computing community who want to modernise legacy Fortran
numerics into differentiable, GPU-ready JAX code.

---

## Features

- **Static analysis** — parse any Fortran codebase, extract modules/subroutines/types,
  build dependency graphs, decompose into translation units.
- **LLM-powered translation** — translates Fortran to JAX with strict differentiability
  rules (no Python control-flow inside jitted functions).
- **Automated tests** — generates `pytest` files for every translated module.
- **Repair loop** — iteratively fixes test failures using root-cause analysis.
- **Unified CLI** — two commands cover the full workflow: `transjax` and `transjax-analyze`.

---

## Installation

```bash
pip install transjax
```

For development / running the test suite:

```bash
git clone https://github.com/AyaLahlou/TransJAX.git
cd TransJAX
pip install -e ".[dev]"
```

---

## Quick start

### 1. Set your Anthropic API key or Claude OAuth token

```bash
export ANTHROPIC_API_KEY="sk-ant-..."
# or copy the template and fill it in
transjax init && cp .env.template .env
```
Pro and Max subscribers can generate an OAuth token by running ```claude setup-token``` locally, which produces a ```CLAUDE_CODE_OAUTH_TOKEN```. This token can be used in place of ```ANTHROPIC_API_KEY``` specifically in Claude Code and the ```claude-code-action``` GitHub Action.
### 2. Inspect the codebase first (optional)

```bash
transjax analyze /path/to/fortran_code
```

### 3. Translate

```bash
transjax convert /path/to/fortran_code -o ./jax_output
```

---

## CLI reference

All functionality lives under the single `transjax` command:

```
Usage: transjax [OPTIONS] COMMAND [ARGS]...

  TransJAX — translate Fortran scientific code to JAX.

  Common workflow:
    transjax analyze /path/to/fortran           # inspect the codebase first
    transjax convert /path/to/fortran -o ./out  # translate + test + repair

Commands:
  analyze      Analyse a Fortran codebase without translating it.
  convert      Translate a Fortran codebase to JAX (full pipeline).
  init         Create a .env.template file in the current directory.
  show-config  Print the active configuration (YAML).
```

#### `transjax convert`

| Option | Default | Description |
|--------|---------|-------------|
| `--output / -o` | `./jax_output` | Output directory |
| `--model` | from config | Claude model name |
| `--api-key` | `$ANTHROPIC_API_KEY` | Anthropic API key |
| `--max-repair-iterations` | 5 | Repair loop limit per module |
| `--skip-tests` | false | Skip test generation |
| `--skip-repair` | false | Skip repair loop |
| `--force` | false | Re-translate existing files |
| `--modules` | all | Comma-separated module filter |
| `--temperature` | 0.0 | LLM temperature |
| `--verbose / -v` | false | Verbose logging |

#### `transjax analyze`

| Option | Default | Description |
|--------|---------|-------------|
| `--output / -o` | `<src>/transjax_analysis` | Output directory |
| `--template / -t` | `auto` | Project template (auto, ctsm, scientific_computing, generic, …) |
| `--no-graphs` | false | Skip graph visualisation |
| `--verbose / -v` | false | Verbose logging |

---

## Python API

```python
from transjax import FortranAnalyzer, create_analyzer_for_project

# Analyse a Fortran project
analyzer = create_analyzer_for_project("/path/to/fortran", template="auto")
results = analyzer.analyze()
print(f"Found {len(results['modules'])} modules")

# Full translation pipeline
from transjax import OrchestratorAgent
from pathlib import Path

orch = OrchestratorAgent(
    fortran_dir=Path("/path/to/fortran"),
    output_dir=Path("./jax_output"),
    skip_tests=True,
)
summary = orch.run()
```

---

## Output structure

```
jax_output/
├── src/                        # Translated Python/JAX code
│   └── <source_dir>/
│       └── <module>.py
├── tests/                      # Generated pytest files
│   └── <source_dir>/
│       └── test_<module>.py
├── docs/                       # Translation notes (Markdown)
├── reports/                    # Logs, summaries, repair logs
│   └── translation_summary.json
└── static_analysis/            # Fortran analyser output
    ├── analysis_results.json
    └── translation_units.json
```

---

## Development

```bash
# Install with all dev dependencies
pip install -e ".[dev]"

# Run tests
pytest tests/ -v

# Lint
ruff check src/

# Build a wheel
pip install build
python -m build --wheel
```

---

## License

BSD-3-Clause.  See [LICENSE](LICENSE).

## Links

- **Repository**: https://github.com/AyaLahlou/TransJAX
- **Issue tracker**: https://github.com/AyaLahlou/TransJAX/issues
