Skip to content

TrueSerien/optorch

Repository files navigation

Optorch / 光计算可微分模拟框架

English: A NumPy-based differentiable photonic computing framework for hybrid optoelectronic AI acceleration. Pure NumPy, zero PyTorch dependency — ML engineers can train optical neural networks without knowing MZI.

中文: 面向光电混合 AI 加速的、基于 NumPy 的可微分光计算框架。纯 NumPy,零 PyTorch 依赖,让机器学习工程师不需要懂 MZI 就能训练光神经网络。


Project Positioning / 项目定位

Optorch aims to be the NumPy-based differentiable simulation framework for photonic computing — the "PyTorch of optics," but fully self-developed with zero external ML framework dependencies.

中文: Optorch 的目标是成为光计算领域的 NumPy 可微分模拟框架,定位类似"光计算的 PyTorch",但完全自研、零外部 ML 框架依赖。

The photonic computing hardware space already has chip companies like Lightelligence and Lightmatter (analogous to NVIDIA making GPUs), but the ecosystem lacks an open-source, flexible, differentiable software framework for ML researchers.

中文: 光计算硬件领域已有"曦智""光本位"等芯片公司(类似 NVIDIA 做 GPU),但光计算生态缺少一个面向 ML 研究者、开源、灵活、可微分的软件框架。


Core Features / 核心特性

  • Pure NumPy Implementation (纯 NumPy 实现): Zero PyTorch dependency, only requires numpy>=1.21.0
  • GPU Acceleration via CuPy (GPU 加速): Auto-detects CuPy; FFT and matrix ops seamlessly migrate to GPU without code changes
  • Complex OpticalField (复数光场): Fundamental photonic data type carrying wavelength, medium, and sampling spacing
  • Physics-Based Adjoint Autodiff (物理伴随法自动微分): Based on optical reciprocity, not conventional chain rule
  • Custom Gradient Registry API (自定义梯度注册 API): Decorator syntax similar to PyTorch autograd.Function
  • Multi-Output Gradient Merging (多输出节点自动合并): Parent node gradients automatically accumulate from all child paths in branched computation graphs
  • Photonic Primitives (光计算原语): Propagation (angular spectrum), phase modulation, beam splitting/coupling, intensity detection, electro-optic modulation
  • MZI Optical Neural Network Layer (MZI 光神经网络层): N×N unitary transform via Clements decomposition
  • Arbitrary Matrix Auto-Decomposition (任意矩阵自动分解): SVD + Clements decomposition; any real/complex matrix → MZI parameters
  • Diffractive Optical Layer (D2NN) (衍射光学层): 4f frequency-domain mask operation (FFT → learnable mask → IFFT)
  • Phase Mask Layer (相位掩模层): Per-pixel independently learnable spatial phase modulation
  • Physical Noise Modeling (物理噪声建模): Quantization, thermal drift, crosstalk, phase error (noise-aware training)
  • Trainer Framework (Trainer 训练框架): Full training loop, early stopping, checkpoint recovery
  • Learning Rate Schedulers (学习率调度器): StepLR, MultiStepLR, ExponentialLR, CosineAnnealingLR, ReduceLROnPlateau
  • Checkpoint Management (Checkpoint 管理): Auto save/restore model, optimizer, and scheduler state
  • Model Serialization (模型序列化): .npz format for full network architecture and parameters
  • Self-Developed Optimizers (自研优化器): SGD and Adam in pure NumPy
  • Batch Training Support (Batch 训练支持): All layers support batched inputs

Quick Start / 快速开始

Installation / 安装

git clone https://github.com/TrueSerien/optorch.git
cd optorch
pip install -e .

GPU Acceleration (Optional) / GPU 加速(可选)

Optorch supports automatic GPU migration via CuPy without modifying any model code:

中文: Optorch 支持通过 CuPy 自动将计算迁移到 GPU,无需修改任何模型代码。

# Install CuPy (choose based on your CUDA version)
pip install cupy-cuda12x   # CUDA 12.x
# pip install cupy-cuda11x # CUDA 11.x

Usage:

中文: 使用方式:

import optorch

# Create optical field (default on CPU)
field = optorch.OpticalField.uniform_plane(shape=(1024, 1024))

# Migrate to GPU
field_gpu = field.to_gpu()

# Build network and forward (FFT auto-executes on GPU)
model = optorch.OpticalNetwork([optorch.MZILayer(n_modes=1024)])
output = model.forward(field_gpu)

# Check device
print(output.device)  # cuda:0

# Move back to CPU
output_cpu = output.to_cpu()

Automatic backend switching: If the input field's data is a cupy.ndarray, all subsequent operations (FFT, phase modulation, MZI transform) automatically use CuPy on GPU. If CuPy is not installed, it gracefully falls back to NumPy CPU mode.

中文: 自动后端切换:如果输入光场的数据是 cupy.ndarray,所有后续操作(FFT、相位调制、MZI 变换)自动使用 CuPy 在 GPU 上执行。如果未安装 CuPy,自动降级到 NumPy CPU 模式。

Environment variable control:

中文: 环境变量控制:

export OPTORCH_BACKEND=auto   # Auto-detect (default)
export OPTORCH_BACKEND=cupy  # Force GPU
export OPTORCH_BACKEND=numpy # Force CPU

Example: Training an Optical Neural Network / 示例:训练光神经网络

import numpy as np
from optorch import OpticalField, MZILayer, OpticalNetwork, Adam
from optorch.core.primitives import detect

# Build optical neural network
model = OpticalNetwork([
    MZILayer(n_modes=8),
    MZILayer(n_modes=8),
])

# Generate input
field = OpticalField.uniform_plane(shape=(8,), amplitude=1.0)

# Forward propagation
output = model.forward(field)
intensity = detect(output)

# Backward propagation (adjoint method)
grad_field = 2 * output.data  # d(I)/d(Ē) = 2 * E
output_node = model.autograd.computation_graph[-1]
gradients = model.autograd.backward(grad_field, output_node)

# Parameter update
optimizer = Adam(model.get_parameters(), lr=0.01)
optimizer.step(gradients)
model.set_parameters(optimizer.parameters)

Run full examples:

中文: 运行完整示例:

# Noise-aware training comparison
python examples/train_mzi_onn.py

# Multi-layer mixed training (PhaseMask + MZI)
python examples/train_mixed_network.py

Architecture / 架构设计

┌──────────────────────────────────────────────────────┐
│  User Interface / 用户接口                             │
│  import optorch                                      │
│  net = OpticalNetwork([MZILayer(...), ...])          │
├──────────────────────────────────────────────────────┤
│  Photonic Core (Pure NumPy) / 光计算核心层             │
│  ┌────────────────────────────────────────────────┐  │
│  │ primitives: propagate / phase_shift /          │  │
│  │             beam_split / detect / modulate     │  │
│  │ autograd: Physics-based adjoint engine       │  │
│  │          GradientRegistry / record_multi_output│  │
│  │ field: OpticalField / Medium                   │  │
│  │ mzi_utils: Unitary matrix / MZI mesh utils   │  │
│  └────────────────────────────────────────────────┘  │
├──────────────────────────────────────────────────────┤
│  Network Layers & Optimizers / 网络层 & 优化器         │
│  ┌────────────────────────────────────────────────┐  │
│  │ MZILayer: N×N unitary (Clements decomposition) │  │
│  │ PhaseMaskLayer: Per-pixel phase modulation   │  │
│  │ DiffractiveLayer: 4f frequency D2NN layer      │  │
│  │ ReshapeLayer: Field dimension transform (1D↔2D)│  │
│  │ OpticalNonlinearity: Optical nonlinear activation│ │
│  │ ConvOpticalLayer: FFT spatial convolution    │  │
│  │ OpticalNetwork: Multi-layer + save/load      │  │
│  │ Trainer: Training loop + early stopping      │  │
│  │ CheckpointManager: Model save/restore        │  │
│  │ decomposition: SVD + Clements matrix factor  │  │
│  │ optimizer: SGD / Adam                          │  │
│  │ schedulers: Step / MultiStep / Exponential /  │  │
│  │             CosineAnnealing / ReduceLROnPlateau│  │
│  └────────────────────────────────────────────────┘  │
├──────────────────────────────────────────────────────┤
│  Physical Noise Modeling / 物理噪声建模                  │
│  ┌────────────────────────────────────────────────┐  │
│  │ PhaseErrorNoise / ThermalDriftNoise           │  │
│  │ CrosstalkNoise / QuantizationNoise            │  │
│  └────────────────────────────────────────────────┘  │
└──────────────────────────────────────────────────────┘

Core Modules / 核心模块

optorch.nn.MZILayer

N×N unitary transform layer based on MZI mesh. Uses Clements rectangular decomposition; N modes require N×(N-1)/2 MZIs.

中文: 基于 MZI mesh 的 N×N 酉变换层。使用 Clements 矩形分解,N 个模式需要 N×(N-1)/2 个 MZI。

layer = MZILayer(n_modes=8, mesh_type="rectangular",
                 phase_noise_std=0.0, quantize_bits=0)
output_field = layer.forward(input_field)

Parameters / 参数:

  • n_modes: Number of optical modes / 光模式数量
  • mesh_type: "rectangular" (Clements) or "triangular" (Reck) / "rectangular" (Clements) 或 "triangular" (Reck)
  • phase_noise_std: Phase noise std during training (noise-aware training) / 训练时相位噪声标准差(噪声感知训练)
  • quantize_bits: Output intensity quantization bits / 输出强度量化位数

optorch.nn.PhaseMaskLayer

Per-pixel learnable phase modulation layer. Applies E_out = E_in * exp(j·φ) independently at each spatial position.

中文: 逐像素可学习相位调制层。对光场每个空间位置独立施加 E_out = E_in * exp(j·φ)

layer = PhaseMaskLayer(shape=(16,))       # 1D
layer = PhaseMaskLayer(shape=(8, 8))      # 2D
output_field = layer.forward(input_field)

optorch.nn.DiffractiveLayer

Diffractive optical layer (D2NN-style). Based on 4f system: input → FFT → frequency mask → IFFT → output.

中文: 衍射光学层(D2NN 风格)。基于 4f 系统:输入 → FFT → 频域掩模 → IFFT → 输出。

# Pure phase mode (unitary, energy-conserving)
layer = DiffractiveLayer(height=64, width=64, mode="phase")

# Complex mode (learnable amplitude + phase)
layer = DiffractiveLayer(height=64, width=64, mode="complex")

output_field = layer.forward(input_field)

optorch.OpticalNetwork

Multi-layer composition + model serialization.

中文: 多层组合 + 模型序列化。

# Build
net = OpticalNetwork([MZILayer(16), PhaseMaskLayer(shape=(16,)), MZILayer(16)])

# Forward
output = net.forward(input_field)

# Save / Load
net.save("model.npz")
net_loaded = OpticalNetwork.load("model.npz")

optorch.optim

Self-developed optimizers in pure NumPy:

中文: 自研优化器(纯 NumPy):

from optorch import SGD, Adam

optimizer = Adam(model.get_parameters(), lr=0.01)
optimizer.step(gradients)
model.set_parameters(optimizer.parameters)

Learning Rate Schedulers / 学习率调度器

6 schedulers with chain composition support:

中文: 6 种调度器,支持链式组合:

from optorch.nn.trainer import StepLR, MultiStepLR, CosineAnnealingLR

scheduler = StepLR(optimizer, step_size=10, gamma=0.5)
# Or:
scheduler = CosineAnnealingLR(optimizer, T_max=30, eta_min=1e-6)

for epoch in range(30):
    train(...)
    scheduler.step()

Arbitrary Matrix Decomposition / 任意矩阵自动分解

SVD + Clements decomposition. Automatically converts any real/complex matrix to MZI parameters:

中文: SVD + Clements 分解,自动将任意实数/复数矩阵转换为 MZI 参数:

from optorch.nn.decomposition import decompose_arbitrary, reconstruct_svd

# Auto-decompose: A = U @ S @ V^H
result = decompose_arbitrary(matrix)  # shape: (m, n), arbitrary m, n
theta, phi = result.U_theta, result.U_phi  # MZI parameters
S, Vh_theta, Vh_phi = result.S, result.Vh_theta, result.Vh_phi

# Reconstruct
A_reconstructed = reconstruct_svd(result)
assert np.allclose(A, A_reconstructed, atol=1e-10)

Trainer Framework / Trainer 训练框架

Full training loop + checkpoint management:

中文: 完整的训练循环 + checkpoint 管理:

from optorch.nn.trainer import Trainer, CheckpointManager

trainer = Trainer(model, optimizer, loss_fn=mse_loss, scheduler=scheduler)

# Train
history = trainer.fit(
    train_data, epochs=50, val_data=val_data,
    early_stop_patience=10, checkpoint_interval=5,
)

# Resume from checkpoint
trainer2 = Trainer(model, optimizer, loss_fn=mse_loss)
trainer2.fit(train_data, epochs=50, resume_from="checkpoints/epoch_10.pt")

Custom Gradient Registry API / 自定义梯度注册 API

Similar to PyTorch autograd.Function, with decorator syntax:

中文: 类似 PyTorch autograd.Function,支持装饰器语法:

from optorch import OpticalAutograd
from optorch.core.autograd import GradientRegistry

# Decorator registration
@GradientRegistry.register("my_op")
def my_backward(output_field, adjoint, params, parents):
    grad = {"param": np.sum(np.real(np.conj(adjoint) * output_field.data))}
    return grad, [adjoint * np.exp(-1j)]

# Direct registration
GradientRegistry.register("other_op", other_backward)

# record() auto-looks up backward_fn from registry
node = autograd.record(output_field, op_name="my_op", params=...)

Multi-Output Gradient Merging / 多输出节点自动合并

Parent node gradients automatically accumulate from all child paths in branched computation graphs:

中文: 分支计算图中父节点梯度自动从所有子路径累加:

# Multi-output ops (e.g., beam_split) auto-merge gradients
out1, out2 = beam_split(field1, field2, theta, phi)
nodes = autograd.record_multi_output(
    [out1, out2],
    op_name="beam_split",
    backward_fns=[backward_out1, backward_out2],
    parents=[p1, p2],
)
# During backward, parent p1/p2 adjoint fields auto-accumulate
# contributions from both out1 and out2

中文: 反向传播时,父节点 p1/p2 的伴随场自动累加 out1 和 out2 的贡献。

Noise Models / 噪声模型

Four physical noise models designed for noise-aware training:

中文: 支持四种物理噪声模型,专为噪声感知训练设计:

from optorch.noise.noise_models import (
    QuantizationNoise,    # DAC quantization / DAC 量化
    PhaseErrorNoise,      # Phase control error / 相位控制误差
    ThermalDriftNoise,    # Thermal drift / 热漂移
    CrosstalkNoise,       # Channel crosstalk / 通道串扰
)

noise = PhaseErrorNoise(error_std=0.05)
x_noisy = noise(x, training=True)   # Inject noise during training / 训练时注入噪声
x_clean = noise(x, training=False)  # No noise during inference / 推理时无噪声

Physics-Based Adjoint Autodiff / 物理伴随法自动微分

Optorch uses physics-based adjoint autodifferentiation via Wirtinger calculus, not the conventional chain rule.

中文: Optorch 使用基于 Wirtinger 微分的物理伴随法自动微分,而非传统的链式法则。

Core formulas:

中文: 核心公式:

Operation / 操作 Forward / 前向 Adjoint Backward / 伴随反向
Phase modulation / 相位调制 E_out = E_in·exp(jφ) dL/dφ = -Im(conj(adj)·E_out)
MZI beam split / MZI 分束 [out1; out2] = U(θ,φ)·[in1; in2] adj_in = U^H·adj_out
Diffraction (FFT) / 衍射 (FFT) Z = FFT(E_in) adj_input = N·IFFT(adj_Z̄)
Diffraction (IFFT) / 衍射 (IFFT) E_out = IFFT(Y) adj_Y = (1/N)·FFT(adj_output)
Intensity detection / 光强检测 I = |E|² Terminal, phase info lost

Comparison / 与现有工具对比

Project Positioning Core Limitation
Photontorch Photonic circuit time/frequency domain simulation Academic license, device-design oriented, not ML
Neurophox MZI mesh optical neural networks Unitary-only, TensorFlow-bound, no hardware backend
Simphony PIC circuit SPICE-level simulation Pure circuit simulation, no ML integration
Optorch ML framework (our position) Pure NumPy, physics adjoint, multi-optical layers, trainable models

中文:

项目 定位 核心问题
Photontorch 光子电路时域/频域模拟 学术许可限制、面向器件设计而非 ML
Neurophox MZI mesh 光神经网络 仅支持酉矩阵、TensorFlow 绑定、无硬件后端
Simphony PIC 电路 SPICE 级仿真 纯电路模拟,无 ML 集成
Optorch ML 框架(我们的位置) 纯 NumPy、物理伴随法、多光学层、模型可训练

Roadmap / 路线图

Phase Status Milestone
Phase 1 Pure NumPy core: OpticalField, OpticalAutograd, photonic primitives, MZILayer, SGD/Adam
Phase 2 Batch support, topological sorting (branched graphs), PhaseMaskLayer, DiffractiveLayer
Phase 3 Noise-aware training, physical noise models, model save/load, multi-layer mixed training examples
Phase 4 Arbitrary matrix MZI auto-decomposition (SVD + Clements), Trainer + Checkpoint + LR schedulers, custom gradient registry API + multi-output auto-merge
Phase 5 🔄 Multi-hardware backend abstraction, compiler layer, real-time photonic chip interface
Phase 6 📋 Photonic computing Model Zoo, community ecosystem, benchmark suite

中文:

阶段 状态 里程碑
Phase 1 纯 NumPy 核心:OpticalFieldOpticalAutograd、光计算原语、MZILayer、SGD/Adam
Phase 2 Batch 支持、拓扑排序(分支图)、PhaseMaskLayerDiffractiveLayer
Phase 3 噪声感知训练、物理噪声模型、模型 save/load、多层混合训练示例
Phase 4 任意矩阵 MZI 自动分解(SVD + Clements)、Trainer + Checkpoint + LR 调度器、自定义梯度注册 API + 多输出自动合并
Phase 5 🔄 多硬件后端抽象、编译器层、实时光芯片接口
Phase 6 📋 光计算 Model Zoo、社区生态、基准测试套件

Project Structure / 项目结构

optorch/
├── optorch/
│   ├── __init__.py           # Top-level exports / 顶层导出
│   ├── core/
│   │   ├── field.py          # OpticalField / Medium
│   │   ├── autograd.py       # OpticalAutograd / Node / GradientRegistry / mse_loss
│   │   ├── primitives.py     # Photonic primitives (with adjoint backward) / 光计算原语(含伴随反向)
│   │   └── mzi_utils.py      # MZI utility functions / MZI 工具函数
│   ├── nn/
│   │   ├── optical_network.py # MZILayer / PhaseMaskLayer / OpticalNetwork
│   │   ├── diffractive_layer.py # DiffractiveLayer
│   │   ├── reshape_layer.py  # ReshapeLayer (1D↔2D)
│   │   ├── nonlinearity.py   # OpticalNonlinearity
│   │   ├── conv_layer.py     # ConvOpticalLayer (FFT convolution)
│   │   ├── decomposition.py  # SVD + Clements matrix decomposition
│   │   └── trainer.py        # Trainer / CheckpointManager / LR schedulers
│   ├── noise/
│   │   └── noise_models.py   # Four physical noise models / 四种物理噪声模型
│   └── optim/
│       └── optimizer.py      # SGD / Adam
├── examples/
│   ├── train_mzi_onn.py       # Noise-aware training comparison / 噪声感知训练对比
│   └── train_mixed_network.py # Multi-layer mixed training example / 多层混合训练示例
├── tests/
│   ├── test_core.py           # Core functionality tests / 核心功能测试
│   ├── test_gradient.py       # Gradient correctness verification / 梯度正确性验证
│   └── test_autograd_api.py   # Custom gradient API tests / 自定义梯度 API 测试
├── setup.py
└── README.md

License / 许可证

Apache 2.0


English: The photonic computing field needs its own software stack. Optorch's opportunity is to become that tool.

中文: 光计算领域需要自己的软件栈。Optorch 的机会是成为那个工具。