Debangshu Mukherjee

Tools & open source

Scientific software

Physical models for simulating measurements and recovering material and instrument parameters. I build tools for optics, electron scattering, RHEED, and photoemission, alongside microscopy analysis software.

GitHub

Selected packages

Click any figure to view it full size.
01 / PACKAGE

Janssen

Optical microscopy & phase retrieval

I built Janssen around modularity: small, reusable optical units that I can string together to model a wide variety of optical setups. The same composition supports forward simulation and inversion, so I can fit physical parameters to measurements without writing a separate reconstruction model.

Gaussian beams with identical intensity and different phase profiles.
Gaussian beams with identical intensity and different phase profiles. Source ↗
Propagation of collimated, converging, and diverging beams.
Propagation of collimated, converging, and diverging beams. Source ↗
Zernike modes describing common optical aberrations.
Zernike modes describing common optical aberrations. Source ↗

My contribution: I built a modular wave-optics package that reuses optical components in forward simulation and inversion.

Python / JAXWave opticsPtychography
GitHub ↗Documentation ↗
Architecture & differentiation

Architecture

  • Separate optical components. Propagation, lenses, apertures, and microscope models live in separate modules. Functions transform a wavefront into a new wavefront, so optical systems can be assembled and their components reused in reconstruction.
  • Fields retain amplitude and phase. OpticalWavefront carries the complex field together with wavelength and sampling information. Keeping phase available is essential for interference, focusing, and phase retrieval.
  • Partial coherence is explicit. CoherentModeSet represents illumination as weighted modes. The coherence and reconstruction modules support partially coherent sources, so the source model can be included in the imaging problem.

JAX & differentiation

  • Data structures for JAX. Wavefronts and mode sets are PyTrees: containers whose numerical arrays can pass through JAX transformations together. Validated factories and array-shape annotations help catch incompatible inputs.
  • Reuse the model for fitting. Automatic differentiation connects a measurement residual to the continuous parameters of the assembled optical setup. Jacobian-free Gauss–Newton with Levenberg–Marquardt damping supports inversion of that same model, using JAX derivatives and conjugate-gradient solves.
  • Compilation and batching. Compatible kernels use jax.jit for compilation and jax.vmap for batches or modes. Grid sizes and other shape-defining choices must stay static; setup and numerical execution have different roles.

Design notes: Architecture ↗ · PyTrees ↗

02 / PACKAGE

Ptyrodactyl

Differentiable electron scattering

With Ptyrodactyl, I model how electrons scatter through a specimen and how that scattering appears at the detector. The package covers multislice propagation, convergent-beam electron diffraction (CBED), 4D-STEM, and reconstruction, with JAX derivatives for fitting specimen and microscope parameters.

Examples from an archived Bi₂Se₃ tutorial.

Top and side views of the Bi₂Se₃ specimen.
Top and side views of the Bi₂Se₃ specimen. Source ↗
Projected electrostatic potential, summed over specimen slices.
Projected electrostatic potential, summed over specimen slices. Source ↗
Simulated Bi₂Se₃ CBED pattern on a logarithmic intensity scale.
Simulated Bi₂Se₃ CBED pattern on a logarithmic intensity scale. Source ↗

My contribution: I develop electron-scattering and reconstruction workflows with JAX derivatives for supported specimen and microscope parameters.

Python / JAXElectron microscopy4D-STEM
GitHub ↗Documentation ↗
Architecture & differentiation

Architecture

  • Separate scattering models. Multislice, Bloch-wave, Fourier–Galerkin, and convergent-Born components have distinct modules. Their physical approximations remain explicit, while workflow functions compose them into simulation and reconstruction tasks.
  • Preserve amplitudes until detection. Scattering models keep complex amplitudes available until an explicit detector operation. Coherent and incoherent distributions are represented separately, making the point where intensities are averaged clear.
  • Shared physical data structures. Equinox data containers represent potentials, beams, and distributions; Potential3D provides a common volumetric potential in volts. Validated constructors establish consistent units, shapes, and numerical types at module boundaries.

JAX & differentiation

  • Optimize declared physical parameters. Differentiable scattering kernels let inverse methods vary supported specimen and microscope parameters. A shared jacobian module provides derivative and sensitivity operations across model families.
  • Keep I/O outside the numerical core. Crystal parsing and HDF5 serialization live in inout, while JAX kernels operate on arrays. This separates file handling from the computation that JAX traces for differentiation and compilation.
  • Support is defined per operation. Use jax.grad, jax.jit, and jax.vmap where the public API guarantees them. The chosen solver, detector reduction, and static model configuration determine the supported transformation path.

Design notes: Architecture & JAX support ↗

03 / PACKAGE

Rheedium

Surface structure through electron diffraction

I develop Rheedium to simulate reflection high-energy electron diffraction (RHEED) and fit its patterns to measured data. It follows the experiment from a crystal surface and grazing electron beam to the detector, with derivatives for adjusting crystal and instrument parameters.

Simulated SrTiO₃, MgO, and Bi₂Se₃ detector patterns at logarithmic intensity.
Simulated SrTiO₃, MgO, and Bi₂Se₃ detector patterns at logarithmic intensity. Source ↗
The grazing-incidence geometry used in RHEED.
The grazing-incidence geometry used in RHEED. Source ↗
Ewald-sphere construction and reciprocal-space rods.
Ewald-sphere construction and reciprocal-space rods. Source ↗

My contribution: I built RHEED simulation and fitting workflows, with an agent interface for discovering tasks and inspecting structured results.

Python / JAXRHEEDSurface science
GitHub ↗Documentation ↗
Architecture & differentiation

Architecture

  • A pipeline that follows the experiment. Crystal parsing and surface preparation feed reciprocal-space geometry, scattering intensities, and detector rendering. Separate modules make structure, surface, and instrument effects easier to inspect and combine.
  • One forward model for simulation and fitting. ReconProblem bundles a forward function, measured data, and a loss or residual. Parameter transforms enforce constraints such as positive thickness or normalized weights before evaluating the physical model.
  • Validation is part of the design. Typed PyTrees carry structures, beams, and patterns. A dedicated audit module checks physical invariants and comparisons against reference data, alongside tests for individual numerical components.

JAX & differentiation

  • Gradients through continuous parameters. Supported paths propagate derivatives back to atom positions, lattice parameters, beam energy, and orientation weights. These sensitivities support fitting a simulated diffraction pattern to experimental data.
  • Compile fixed-shape work. Form-factor, potential, and structure-factor kernels with fixed input shapes support JAX compilation and batching. Model-selection strings must be static, and repeated inversions can reuse a persistent compilation cache.
  • Geometry can change array sizes. Top-level builders that choose grid dimensions or select a variable number of reflections cannot be directly compiled as written. Build those structures outside the compiled loop, or fix and pad their sizes; gradients apply within the chosen numerical representation.

Design notes: JAX transformation guide ↗ · Inverse problems ↗

04 / PACKAGE

diffpes

From electronic structure to photoemission

In diffpes, I connect electronic structure calculations with angle-resolved photoemission spectroscopy (ARPES). Starting from tight-binding or density-functional-theory models, the package calculates spectra, Fermi surfaces, and detector signals, including polarization, broadening, and instrument effects.

Graphene band structure and the corresponding simulated ARPES spectrum.
Graphene band structure and the corresponding simulated ARPES spectrum. Source ↗
A three-dimensional ARPES intensity cube around a graphene Dirac point.
A three-dimensional ARPES intensity cube around a graphene Dirac point. Source ↗
An intrinsic spectrum compared with the signal after instrument effects.
An intrinsic spectrum compared with the signal after instrument effects. Source ↗

My contribution: I built electronic-structure-to-photoemission simulation and agent interfaces for compiled scientific workflows.

Python / JAXARPESElectronic structure
GitHub ↗Documentation ↗
Architecture & differentiation

Architecture

  • Separate electronic structure from measurement. Band models and imported electronic structure feed spectral and photoemission calculations, followed by instrument effects. This lets the physical spectrum and the measured detector signal be examined separately.
  • Preserve quantum-mechanical phase. Complex band coefficients remain available to coherent matrix-element calculations. Orbital projection probabilities are stored separately because they do not contain the phase information needed to reconstruct interference.
  • Distinguish parameters from model structure. Equinox PyTrees hold numerical values such as geometry, radial coefficients, and temperature. Orbital quantum numbers, basis mappings, and model selectors are static metadata; factories validate both before numerical execution.

JAX & differentiation

  • Differentiate spectra and matrix elements. Numerical kernels support automatic differentiation, compilation, and vectorization for continuous inputs. Parameter-packing helpers expose active matrix-element variables as a real vector for optimization while retaining fixed calibration data.
  • Handle degeneracy explicitly. A resolvent-based spectral path avoids directly differentiating individual eigenvectors at degeneracy. Sensitivity tools can operate on complete isolated band groups, whose combined response has a meaningful physical interpretation.
  • Expose unidentifiable and invalid directions. Gauge helpers identify parameter changes, such as a common phase, that leave intensity unchanged. Validity masks mark dark or invalid regions; optimizers must respect these boundaries. File parsing and static basis selection remain outside the differentiated numerical path.

Design notes: PyTree architecture ↗ · Gradients & boundaries ↗

Agent access to scientific workflows

I built end-to-end RHEED and ARPES simulators and integrated chatbots through their automatons/ workflows. These expose ahead-of-time (AOT) compiled JAX computations as discoverable tasks, so an agent can use the simulator without reading thousands of lines of scientific code.

The agent selects tasks and reads structured results; the JAX workflows carry out the simulation and reconstruction. These interfaces make the scientific tools accessible through a chatbot.

Rheedium agent workflows ↗diffpes agent workflows ↗
Technical background: why I use JAX

I use the same physical model to simulate a measurement and fit it to an experiment. JAX lets me calculate how the predicted signal changes with the model parameters, so an optimizer can use those derivatives to improve the fit.

  • Automatic differentiation: jax.grad computes derivatives of scalar losses through compatible numerical operations, supporting parameter fitting and sensitivity analysis.
  • Compilation: jax.jit compiles compatible functions for efficient execution on available hardware, including CPUs and GPUs.
  • Batching: jax.vmap maps a function over a batch dimension, helping express parameter sweeps or ensembles without a Python loop for every sample.
  • Structured parameters: PyTrees group arrays into physical objects, so model parameters can be passed through these transformations together.

Each package defines which numerical paths support these transformations. Differentiation concerns continuous parameters; changing a grid size, model choice, or discrete structure requires separate treatment. JAX concepts ↗

More of my software

The packages above are a selection of my work. I also build tools for detecting features and analyzing experimental microscopy data.

CryoBlob

I built CryoBlob for reference-free detection of compact features in low-signal cryo-EM images, using JAX compilation and GPU acceleration.

GitHub ↗Documentation ↗

stemtool

I built stemtool as an open source Python toolkit for analyzing scanning transmission electron microscopy datasets.

GitHub ↗Documentation ↗