Proof-of-concept: Using ChatGPT to Translate and Modernize an Earth System Model from Fortran to Python/JAX (2405.00018v1)

Abstract: Earth system models (ESMs) are vital for understanding past, present, and future climate, but they suffer from legacy technical infrastructure. ESMs are primarily implemented in Fortran, a language that poses a high barrier of entry for early career scientists and lacks a GPU runtime, which has become essential for continued advancement as GPU power increases and CPU scaling slows. Fortran also lacks differentiability - the capacity to differentiate through numerical code - which enables hybrid models that integrate machine learning methods. Converting an ESM from Fortran to Python/JAX could resolve these issues. This work presents a semi-automated method for translating individual model components from Fortran to Python/JAX using a LLM (GPT-4). By translating the photosynthesis model from the Community Earth System Model (CESM), we demonstrate that the Python/JAX version results in up to 100x faster runtimes using GPU parallelization, and enables parameter estimation via automatic differentiation. The Python code is also easy to read and run and could be used by instructors in the classroom. This work illustrates a path towards the ultimate goal of making climate models fast, inclusive, and differentiable.

• The paper explores using ChatGPT/LLMs as a proof-of-concept to semi-automate the translation of Earth System Models from legacy Fortran code to modern Python/JAX.
• A case study on a photosynthesis module demonstrated significant performance improvements, achieving up to 100x speedup on GPUs, and enabled automatic differentiation for efficient parameter optimization.
• This approach aims to improve the accessibility of climate modeling for researchers by using Python and facilitates future integration with machine learning techniques for enhanced model accuracy.

Essay: Utilizing ChatGPT for Translating Earth System Models from Fortran to Python/JAX

The paper "Proof-of-concept: Using ChatGPT to Translate and Modernize an Earth System Model from Fortran to Python/JAX" investigates the feasibility of using LLMs, specifically GPT-4, as a tool to modernize Earth System Models (ESMs). ESMs, pivotal in climate science, have traditionally been developed in Fortran, a language that imposes certain technical barriers and inefficiencies, such as limited differentiability and poor adaptability to GPUs. This research presents an approach to transition these models to Python/JAX, with the goal of harnessing modern computational capabilities such as GPU acceleration and automatic differentiation, thereby bolstering accessibility and performance.

Overview of the Translation Methodology

The methodology is centered around a semi-automated translation approach, leveraging GPT-4 to convert Fortran segments to Python/JAX. The translation process relies on a divide-and-conquer strategy, where the Fortran codebase is partitioned into smaller, manageable units that are comprehensible to the LLM, due to inherent context length limitations. This approach includes:

1. Static Analysis and Dependency Ordering: By employing static analysis, the researchers delineate the codebase into discrete units. A topological sort of these units, based on dependencies, ensures correct translation sequence.
2. Iterative Code Generation and Testing: Each unit is iteratively translated and refined using GPT-4 until the generated Python code successfully passes a comprehensive suite of unit tests. This iterative process helps overcome potential inaccuracies in initial LLM outputs.

Evaluation and Results

Significantly, the translation of a leaf-level photosynthesis module from the Community Earth System Model (CESM) is presented as a case paper. The paper reports substantial improvements in computational efficiency, with the Python/JAX implementation achieving up to 100x speedup on GPUs compared to its Fortran counterpart on CPU. This dramatic performance enhancement underscores the potential of modern hardware utilization in climate modeling.

Furthermore, the inclusion of automatic differentiation through JAX facilitates efficient parameter estimation. The paper demonstrates this advantage by optimizing photosynthesis-related parameters using gradient descent, a feat impractical in the original Fortran framework. The ability to perform such optimizations opens avenues for refined model tuning and enhanced precision in simulations.

Implications and Future Directions

The implications of this work are profound, particularly in rendering ESMs more accessible to early-career scientists unfamiliar with legacy languages like Fortran. By adopting Python, which is widely used across scientific domains, the entry barrier decreases, enabling broader participation and innovation.

Theoretically, the migration to Python/JAX positions climate models to exploit machine learning advancements. The capacity for real-time model updates through online learning and integration with neural network-based subgrid processes could significantly enhance model accuracy and predictive capabilities.

Future research may address scaling the translation process to encompass full climate models. Challenges, such as Fortran's complex module interdependencies and GPT-4's token limitations, are non-trivial and require innovative solutions. Potential advancements include leveraging more sophisticated compiler representations or integrating logging mechanisms to facilitate more seamless translation.

Conclusion

The paper provides valuable insights into modernizing the computational infrastructure of ESMs. The semi-automated translation method not only demonstrates a feasible path to leveraging advanced computational tools like Python/JAX but also sets the foundation for making climate models faster, more accurate, and inclusive. As climate change demands increasingly sophisticated modeling strategies, the transition to adaptive, high-level programming languages represents a critical step forward in the scientific community's capacity to simulate and understand Earth's complex systems.