SciJo: Scientific computing library for Mojo

Community Showcase
1

Hello everyone!

I’m excited to announce the first release of SciJo, a high-performance scientific computing library for Mojo!

What is SciJo?

SciJo brings SciPy-like functionality to Mojo, with native performance and strong type safety.

I’ve been using Mojo extensively in my research, and over time, I began consolidating the utilities I developed into a single, cohesive library. That effort evolved into SciJo, an open project with the ambitious goal of bringing the full range of SciPy capabilities to Mojo.

SciJo builds on NuMojo for array operations, making NuMojo and SciJo together a powerful foundation for scientific computing in Mojo.

While we’re just getting started, v0.1 already includes several core modules ready to use today.

What’s Included in v0.1

Numerical Differentiation

  • Central, forward, and backward finite differences
  • Adaptive step size refinement with Richardson extrapolation
  • Multiple accuracy orders (1-6 for forward/backward, 2-8 for central)

∫ Integration

  • quad: Gauss-Kronrod quadrature (QUADPACK QNG algorithm)
  • trapezoid: Classic trapezoidal rule for uniform/non-uniform grids

Interpolation

  • interp1d: Linear interpolation with extrapolation support
  • Fully compatible with NuMojo arrays

FFT

  • fft/ifft: Fast Fourier Transform using Cooley-Tukey algorithm
  • NumPy-compatible conventions for easy migration
  • Support for complex arrays (power-of-2 sizes)

Physical Constants

  • CODATA 2022 fundamental constants
  • Values, units, and uncertainties included
  • Drop-in replacement for scipy.constants

Quick Example

import scijo as sj
from scijo.integrate import quad

fn integrand[dtype: DType](x: Scalar[dtype], args: Optional[List[Scalar[dtype]]]) -> Scalar[dtype]:
    return x * x

fn main() raises:
    var result = quad[sj.f64, integrand](a=0.0, b=1.0, args=None, epsabs=1e-6, epsrel=1e-6)
    print("∫x² from 0 to 1 =", result.integral)  # Should be ~0.333

Installation

Via pixi (recommended):
Add the following to your pixi.toml

[workspace]
preview = ["pixi-build"]

[dependencies]
modular = ">=25.6.1,<26"
scijo = { git = "https://github.com/shivasankarka/SciJo.git", branch = "main"}

Full installation instructions: GitHub - shivasankarka/SciJo: A high-performance scientific computing library for Mojo, providing SciPy-like functionality with the speed and efficiency of native Mojo code.

What’s Next

I’m actively developing SciJo and have an exciting roadmap ahead:

Near term:

  • More integration algorithms (QAGSE, Simpson’s, Romberg etc)
  • Add GPU support (NuMojo Backend) in modules where it’s possible such as FFT.
  • More FFT (rfft, irfft) and 2D FFT support
  • Additional interpolation methods (cubic, spline)
  • Expand differentiation module.

Future:

  • Optimization (minimization, root finding)
  • Statistical functions and distributions
  • Signal processing (filtering, convolution)
  • Linear algebra (decompositions)

Get Involved

Contributions are most welcome! Whether you want to:

  • Implement new algorithms
  • Optimize existing code
  • Write tests or documentation
  • Report bugs or suggest features

Feel free to open a PR or issue on GitHub.

Let’s build the scientific computing ecosystem for Mojo together! If you use SciJo in your work, consider citing it, it helps with visibility and supports the project’s growth :slight_smile:

Happy computing!

Links:

11 Likes
2

This is great! I’m happy to see more python staples gradually getting pure mojo versions!

You might find Kelvin very interesting. It does unit management in the type system, which enables users to avoid mistakes like adding meters and frequency by making it into a compiler error. I’d also suggest checking out the standard library more deeply, there’s a lot of stuff in there which can help you, like the newly merged fft implementation and a lot of float accuracy stuff. NuMojo’s datatypes are all re-exported Mojo stdlib types, so you can use everything there.

You might also want to try out some benchmarks. We’ve had more than a few people port a library and then beat the C/C++ code in the original by accident.

2 Likes
3

This is really really cool, I’m thrilled to see this. Great job!

-Chris

1 Like
4

Yes, thanks and good job !

5

I second this. Very happy to see this, but also very curious how the benchmarks look in comparison to numpy.
Thanks for working on this! :smiley:

6

Thanks for sharing! @owenhilyard, I do want to take a look at the fft implementation by Martin since I am very new to GPU stuff and I am currently working on GPU backend for NuMojo. I have been thinking about using Kelvin as a backend for SciJo constants module, but I remember there were some discussions going on related to its compile time being quite long. I will take a peek again. Also, could you refer me to the relevant stdlib stuff you were talking about?

7

We’ve had more than a few people port a library and then beat the C/C++ code in the original by accident.

For what It’s worth, I’m in the middle of porting some C++ code to Mojo, and I also beat the original by accident. It seems Mojo’s explicit vectorization tools can beat g++’s autovectorizer even in really simple cases where autovectorization should work pretty well.

I haven’t looked at the assembly directly yet to see exactly what’s going on there, but after only the dumbest optimization pass over the Mojo code (took out some accidental string building in a debug_assert), the Mojo code was about 2x faster than the equivalent C++ code. It seems Mojo’s design makes the fast thing also the easy thing to do!

2 Likes
8

@samufi We’re currently focused on porting all the core functionalities from NumPy and SciPy into NuMojo and SciJo. At the same time, we’re adding a GPU backend to NuMojo, so we haven’t spent much time on benchmarking yet. I’m also still figuring out the best approach to benchmark NumPy for a fair comparison with NuMojo. But I am super excited to sit down, do some benchmarking and get down the rabbit hole of optimization! :slight_smile:

1 Like
9

As another scientific user of Mojo, I’m also very interested in FFTs. If you’re looking for inspiration (or a challenge), we’re interested in real-to-complex and complex-to-real 2d and 3d transforms, where the data dimensions are only known at runtime.

These are pretty challenging conditions for FFT, I know, but I had some ideas where Mojo could make a unique contribution to this space.

We’re currently using the fftw library to handle transforms (on the CPU only for now, sadly). fftw handles dynamic data sizes by using a planning phase to decide things like vector sizes, adjustments due to memory alignments, tree structure for the recursion, etc.

It seems to me these planing steps are all things a compiler should be able to mostly do, perhaps aided by a little metaprogramming for the tree structure bits. So my idea for FFT in Mojo was to write the FFT code using mostly parameterization, including for data sizes, like Martin’s approach so far. But at compile time, only compile the FFT code to the phase right before parameter substitution. Then at runtime, rely on some kind of JIT to substitute the parameter values (like the data sizes) to get the final assembly. I suspect building the FFT code with a mix of AOT and JIT would result in something even faster than fftw can manage at the moment. And it could potentially be very portable, like all Mojo code is.

I’m not sure if Mojo supports this kind of JIT just yet? But it seems like it’s within the realm of possibility. I’m curious what others think about the idea.

Mixed AOT/JIT tools could be really useful for scientific computing in general, where it’s generally favorable to trade off a little up-front planning time if it can speedup the overall workload significantly. And for many applications, shipping the compiler tools with the app isn’t such a problem.

10

Thank you for saying this - this is exactly what we’re trying to achieve with Mojo, and I’m thrilled it is working well for you!

As another scientific user of Mojo, I’m also very interested in FFTs. If you’re looking for inspiration (or a challenge), we’re interested in real-to-complex and complex-to-real 2d and 3d transforms, where the data dimensions are only known at runtime.

If you are interested, @martinvuyk has done some really impressive work on GPU FFTs, beating cuFFT by quite a lot. He recently presented it at the october community meeting. The video should be up here soon:

I’m not sure if Mojo supports this kind of JIT just yet? But it seems like it’s within the realm of possibility. I’m curious what others think about the idea.

Mixed AOT/JIT tools could be really useful for scientific computing in general, where it’s generally favorable to trade off a little up-front planning time if it can speedup the overall workload significantly. And for many applications, shipping the compiler tools with the app isn’t such a problem.

The answer is “yes” - this is exactly what Mojo+MAX do in conjunction. MAX is an AI framework that takes an “operator graph” (which can be a single operator) and then applies various optimizations to it. These optimizations include things like specialization based on staticly-known information like shapes/dimensions, as well as other things like memory planning and fusion of different kernels to reduce memory traffic

-Chris

2 Likes
11

Thank you for saying this - this is exactly what we’re trying to achieve with Mojo, and I’m thrilled it is working well for you!

I’ve been super impressed with Mojo so far, so nice work! Even though we’re not doing anything remotely related to genAI, Mojo is quickly proving to be the best tool available for general scientific HPC anyway.

If you are interested, @martinvuyk has done some really impressive work on GPU FFTs, beating cuFFT by quite a lot.

Yes! I’ve been discussing with Martin about how I can use their new FFT code. Unfortunately, the dynamic sizes issue has been a blocker for my work, but I’m hoping we can find a solution (that doesn’t sacrifice performance) by exploring some kind of JIT approach. It’s encouraging to hear that MAX already does something similar, so maybe there’s hope we can get it working without too much trouble.

I think much of the MAX API was open-sourced recently? So I’ll take a look and try to get an idea of how it works. Maybe using the graph API directly is what we need in this case? Or maybe there’s some lower-level tool in there somewhere that’s a better fit since we’re not actually doing anything with AI models here.