Skip to content

Commit

Permalink
works
Browse files Browse the repository at this point in the history
  • Loading branch information
@rjmoerland
rjmoerland committed Nov 4, 2024
1 parent 87450ac commit 9909df0
Show file tree
Hide file tree
Showing 6 changed files with 352 additions and 1,266 deletions.
3 changes: 0 additions & 3 deletions src/lumicks/pyoptics/trapping/__init__.py
View file
Original file line number Diff line number Diff line change
@@ -1,12 +1,9 @@
from .bead import Bead
from .interface import (
fields_focus,
_fields_focus,
fields_focus_gaussian,
fields_plane_wave,
_fields_plane_wave,
forces_focus,
_forces_focus,
force_factory,
scattered_power_focus,
absorbed_power_focus,
Expand Down
214 changes: 7 additions & 207 deletions src/lumicks/pyoptics/trapping/focused_field_calculation.py
View file
Original file line number Diff line number Diff line change
@@ -1,24 +1,16 @@
from typing import Tuple

import numpy as np
from numba import njit

from ..objective import Objective
from .bead import Bead
from .implementation import (
R_pol_R_th_R_phi,
R_th_R_phi,
internal_field_fixed_r,
internal_H_field_fixed_r,
scattered_field_fixed_r,
scattered_H_field_fixed_r,
)
from .legendre_data import calculate_legendre
from .local_coordinates import (
ExternalBeadCoordinates,
InternalBeadCoordinates,
LocalBeadCoordinates,
)
from .numba_loop import _do_loop
from .radial_data import calculate_external as calculate_external_radial_data
from .radial_data import calculate_internal as calculate_internal_radial_data

Expand Down Expand Up @@ -72,14 +64,16 @@ def calculate_field(
calculate_magnetic_field: bool = False,
calculate_total_field: bool = True,
):
dummy = np.atleast_2d(0)
local_coords = local_coordinates.xyz_stacked
region = np.reshape(local_coordinates.region, local_coordinates.coordinate_shape)

# Always provide a (dummy) 2D numpy array to satisfy Numba (can't deal with None it seems),
# but keep it small in case the parameter isn't used
dummy = np.atleast_2d(0)

# It's ugly but at least Numba compiles the loop
E_field, H_field = _do_loop(
bead_center[0],
bead_center[1],
bead_center[2],
bead_center,
an,
bn,
cn,
Expand Down Expand Up @@ -137,197 +131,3 @@ def calculate_field(
return ret_val

return calculate_field


@njit(cache=True)
def _do_loop(
bead_x,
bead_y,
bead_z,
an,
bn,
cn,
dn,
n_bead,
n_medium,
krH,
dkrH_dkr,
k0r,
sphBessel,
jn_over_k1r,
jn_1,
aperture,
cos_theta,
sin_theta,
cos_phi,
sin_phi,
kx,
ky,
kz,
Einf_theta,
Einf_phi,
legendre_data,
legendre_data_dtheta,
r,
local_coords,
internal: bool,
total: bool,
calculate_electric: bool,
calculate_magnetic: bool,
):
local_cos_theta = np.empty(r.shape)
plane_wave_response_xyz = np.empty((3, r.size), dtype="complex128")
# Mask r == 0:
r_eq_zero = r == 0

dummy = np.zeros((1, 1), dtype="complex128")
field_storage_E = np.zeros_like(plane_wave_response_xyz) if calculate_electric else dummy
field_storage_H = np.zeros_like(plane_wave_response_xyz) if calculate_magnetic else dummy

# Skip points outside aperture
rows, cols = np.nonzero(aperture)
if r.size > 0:
for row, col in zip(rows, cols):
matrices = [
R_th_R_phi(
cos_theta[row, col],
sin_theta[row, col],
cos_phi[row, col],
-sin_phi[row, col],
),
R_pol_R_th_R_phi(
cos_theta[row, col],
sin_theta[row, col],
cos_phi[row, col],
-sin_phi[row, col],
),
]

E0 = [Einf_theta[row, col], Einf_phi[row, col]]

for polarization in range(2):
A = matrices[polarization]
coords = A @ local_coords
x = coords[0, :]
y = coords[1, :]
z = coords[2, :]
if polarization == 0:
local_cos_theta[:] = z / r
local_cos_theta[r_eq_zero] = 1
np.clip(local_cos_theta, a_max=1, a_min=-1, out=local_cos_theta)

# Expand the legendre derivatives from the unique version
# of cos(theta)
local_sin_theta = ((1 + local_cos_theta) * (1 - local_cos_theta)) ** 0.5
indices = legendre_data[1][row, col]
alp_sin_expanded = legendre_data[0][:, indices]
alp_expanded = alp_sin_expanded * local_sin_theta
indices = legendre_data_dtheta[1][row, col]
alp_deriv_expanded = legendre_data_dtheta[0][:, indices]

rho_l = np.hypot(x, y)
cosP = x / rho_l
sinP = y / rho_l
where = rho_l == 0
cosP[where] = 1
sinP[where] = 0

if calculate_electric:
if internal:
plane_wave_response = internal_field_fixed_r(
cn,
dn,
sphBessel,
jn_over_k1r,
jn_1,
alp_expanded,
alp_sin_expanded,
alp_deriv_expanded,
local_cos_theta,
local_sin_theta,
cosP,
sinP,
)
else:
plane_wave_response = scattered_field_fixed_r(
an,
bn,
krH,
dkrH_dkr,
k0r,
alp_expanded,
alp_sin_expanded,
alp_deriv_expanded,
local_cos_theta,
local_sin_theta,
cosP,
sinP,
total,
)

plane_wave_response_xyz[:] = A.T.astype("complex128") @ plane_wave_response
plane_wave_response_xyz *= (
E0[polarization]
* np.exp(
1j
* (
kx[row, col] * bead_x
+ ky[row, col] * bead_y
+ kz[row, col] * bead_z
)
)
/ kz[row, col]
)

field_storage_E += plane_wave_response_xyz

if calculate_magnetic:
if internal:
plane_wave_response = internal_H_field_fixed_r(
cn,
dn,
sphBessel,
jn_over_k1r,
jn_1,
alp_expanded,
alp_sin_expanded,
alp_deriv_expanded,
local_cos_theta,
local_sin_theta,
cosP,
sinP,
n_bead,
)
else:
plane_wave_response = scattered_H_field_fixed_r(
an,
bn,
krH,
dkrH_dkr,
k0r,
alp_expanded,
alp_sin_expanded,
alp_deriv_expanded,
local_cos_theta,
local_sin_theta,
cosP,
sinP,
n_medium,
total,
)
plane_wave_response_xyz[:] = A.T.astype("complex128") @ plane_wave_response
plane_wave_response_xyz *= (
E0[polarization]
* np.exp(
1j
* (
kx[row, col] * bead_x
+ ky[row, col] * bead_y
+ kz[row, col] * bead_z
)
)
/ kz[row, col]
)
# Accumulate the field for this plane wave and polarization
field_storage_H += plane_wave_response_xyz
return field_storage_E, field_storage_H
Loading

0 comments on commit 9909df0

Please sign in to comment.