A different path in machine learning and plasma physics research

Leaving my postdoc at UCLA, I was an academic with a narrow perspective and a middling scientific programmer. Since, I have had the opportunity to work for a few different organizations in a few different capacities with a few different teams, where I was able to grow into a more well rounded researcher, software engineer, and teammate.

While spending my time tinkering with neural networks and related technology at startups and big tech, it became clear to me that the novel technologies surrounding the deep learning revolution had much to offer to plasma physics and fusion.

I wanted to build fast, flexible, and differentiable plasma tools. With the guidance and hard work of mentors and collaborators, we presented, published, and brought this work out of side-project-land to make it my primary focus.

Ergodic LLC is the result. It is a small research organization that both, develops software for accelerating plasma physics, and implements this software in order to uncover physical truths, optimize experiments, and analyze data in novel ways. We do this in collaboration with other academic and research institutions.

At the time of writing, the discussion around deep neural networks with respect to scientific computing applications usually starts and ends with “a faster-function-approximator”. However, there are other more subtle but arguably more important advances. A shared principle among these is Python’s ability to cater to different aspects of the scientific computing process. The three advances that are paramount to Ergodic's success are

1. the improvements in the accessibility and performance of Pythonic Automatic Differentiation packages. While the features provided by these tools are not necessarily new (implementations exist in C++ and FORTRAN), the availability in Python (and Julia?) has driven a rapid increase in adoption due to the flexible and accessible nature of Python.

The result is the ability to apply AD to general purpose numerical programs. This has the potential to unblock branches of plasma physics that have been trapped by problems like analytical intractability and impossibly large computational demands.

2. the improvements in the accessibility and performance of Accelerated Vectorized computations (e.g. CUDA in Python). These improvements directly impact the ability to develop interesting general purpose scientific computing software because Python is finally “fast enough”.

Again, combining Python’s flexibility, iterability, and ecosystem with CUDA-level performance can unblock significant portions of plasma physics.

3. the proliferation of MLOps tools (e.g. MLFlow). This has dramatically improved accessibility to scientific (simulation and real-world) data management platforms.

As is well established in ML/DL, having a sound data strategy, ideally one that is tightly integrated with the model creation and deployment process, is a fundamental component to building good models. There is no reason this does not apply to small groups performing research in the hard sciences. At Ergodic, we have a secured platform, built on AWS, that is accessible via a browser from anywhere in the world, for FAIR [2] access to data. Having flexible, usable, easily deployable, and low maintenance software to help with this goes A LONG WAY.

Realizing that there are shared abstractions at various levels between deep learning, numerical programming, and the experimental sciences enables scientists steeped in numerical programming and physical experiments to implement a data-driven approach that marries the best of all worlds (using the right combination of software and computing tools!)

The equations and approaches we have developed to parse physics can be retained. These serve as hard mathematical constraints, provided by theoretically- and/or empirically- motivated models. We combine these with the modern data-driven modeling techniques to provide added flexibility to the models.

We still write systems of (partial-differential-) equations and solve them via traditional means. However, because our programs are differentiable, we are able to embed and train neural networks into any component of a system of equations, rather than replacing it altogether.

So this is a lower level synergy that leverages, and ideally, enhances the usual approach to plasma physics, rather than replaces it [3]. Because of this, differentiable programming can be applied towards theoretical, computational, and experimental aspects of plasma physics.

Look forward to blog posts and other write-ups here, and datasets and talks and other artifacts on https://ergodic.io

Thanks!

References

[1] - Karpathy - Software 2.0 - https://karpathy.medium.com/software-2-0-a64152b37c35

[2] - FAIR data - https://en.wikipedia.org/wiki/FAIR_data

[3] - P.S. This sounds like The Bitter Lesson - http://www.incompleteideas.net/IncIdeas/BitterLesson.html

The Bitter Lesson states that the best hopes of symbol- and rule- based AI were dashed by a purely data-driven approach once enough compute and data started becoming available.

Does this apply to the sciences / plasma physics? Let’s save that discussion for a future post