Практика с SVETlANNa

pip install svetlanna

import torch
import svetlanna as sv

# Параметры
wavelength = 632.8e-9  # He-Ne лазер, м
pixel_size = 10e-6     # размер пикселя, м
size = 256             # размер сетки
 
# Гауссов пучок
field = sv.GaussianBeam(
    wavelength=wavelength,
    waist=100e-6,
    grid_size=(size, size),
    pixel_size=pixel_size
)

# Распространение методом углового спектра
propagator = sv.AngularSpectrum(
    wavelength=wavelength,
    pixel_size=pixel_size,
    distance=0.1  # 10 см
)
 
output_field = propagator(field)

# Создание обучаемой фазовой маски
phase_mask = sv.PhaseMask(
    grid_size=(size, size),
    init='random'  # случайная инициализация
)
 
# Применение к полю
modulated = phase_mask(field)

# Тонкая линза
lens = sv.ThinLens(
    focal_length=0.05,  # 5 см
    wavelength=wavelength,
    grid_size=(size, size),
    pixel_size=pixel_size
)

# Дифракционная нейросеть
model = sv.D2NN(
    num_layers=5,
    grid_size=(size, size),
    pixel_size=pixel_size,
    wavelength=wavelength,
    layer_distance=0.03  # 3 см между слоями
)

optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
criterion = torch.nn.CrossEntropyLoss()
 
for epoch in range(100):
    optimizer.zero_grad()
 
    # Прямой проход
    output = model(input_field)
    intensity = torch.abs(output) ** 2
 
    # Вычисление потерь
    loss = criterion(intensity, labels)
 
    # Обратный проход
    loss.backward()
    optimizer.step()

import matplotlib.pyplot as plt
 
# Амплитуда и фаза
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4))
 
ax1.imshow(torch.abs(field).cpu())
ax1.set_title('Амплитуда')
 
ax2.imshow(torch.angle(field).cpu(), cmap='twilight')
ax2.set_title('Фаза')
 
plt.show()