torchopticsy 0.8.3__tar.gz

torchopticsy-0.8.3/PKG-INFO

@@ -0,0 +1,26 @@
+Metadata-Version: 2.1
+Name: torchopticsy
+Version: 0.8.3
+Summary: PyTorch-based optics caculation
+Author: YuningYe
+Author-email: 1956860113@qq.com
+License: MIT
+Description-Content-Type: text/markdown
+Requires-Dist: torch
+Requires-Dist: opencv-python
+Requires-Dist: matplotlib
+
+# torchOpticsY
+A PyTorch-based optics calculation library.
+
+
+## Features
+- Luneburg integral 
+- Debye鈥揥olf integral.
+
+## Install
+pip install torchOpticsY
+
+
+0.2 - Interference: OffAxisPhase
+0.1 - Diffraction: Debye_Wolf,Luneburg

torchopticsy-0.8.3/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

torchopticsy-0.8.3/setup.py

@@ -0,0 +1,16 @@
+import pathlib
+import setuptools
+
+setuptools.setup(
+    name="torchopticsy",
+    version="0.8.3",
+    description="PyTorch-based optics caculation",
+    long_description=pathlib.Path("README.md").read_text(),
+    long_description_content_type="text/markdown",
+    author="YuningYe",
+    author_email="1956860113@qq.com",
+    license="MIT",
+    packages=setuptools.find_packages(),
+    install_requires=["torch", "opencv-python", "matplotlib"],
+    include_package_data=True,
+)

torchopticsy-0.8.3/torchopticsy/Diffraction.py

@@ -0,0 +1,680 @@
+import torch
+import torch.nn.functional as F
+
+
+class BluesteinDFT:
+    # Due to the error in periodic expansion, it can be treated by making up for zeros in the future
+    # Now only the central region is relatively accurate, and the surrounding area is affected by pseudo-diffraction by periodic conditions
+    def __init__(self, f1, f2, fs, mout, m_input, device="cpu"):
+        self.device = torch.device(device)
+        self.f1 = f1
+        self.f2 = f2
+        self.fs = fs
+        self.mout = mout
+        self.m_input = m_input
+
+        # Frequency adjustment
+        f11 = f1 + (mout * fs + f2 - f1) / (2 * mout)
+        f22 = f2 + (mout * fs + f2 - f1) / (2 * mout)
+        self.f11 = f11
+        self.f22 = f22
+
+        # Chirp parameters
+        a = torch.exp(1j * 2 * torch.pi * torch.tensor(f11 / fs))
+        w = torch.exp(-1j * 2 * torch.pi * torch.tensor(f22 - f11) / (mout * fs))
+        self.a = a.to(self.device)
+        self.w = w.to(self.device)
+
+        h = torch.arange(
+            -m_input + 1,
+            max(mout - 1, m_input - 1) + 1,
+            device=self.device,
+            dtype=torch.float64,
+        )
+        h = self.w ** ((h**2) / 2)
+
+        self.h = h
+        self.mp = m_input + mout - 1
+        padded_len = 2 ** int(torch.ceil(torch.log2(torch.tensor(self.mp))))
+        self.padded_len = padded_len
+
+        h_inv = torch.zeros(padded_len, dtype=torch.complex64, device=self.device)
+        h_inv[: self.mp] = 1 / h[: self.mp]
+        self.ft = torch.fft.fft(h_inv)
+
+        b_exp = torch.arange(0, m_input, device=self.device)
+        self.b_phase = (self.a**-b_exp) * h[m_input - 1 : 2 * m_input - 1]
+
+        l = torch.linspace(0, mout - 1, mout, device=self.device)
+        l = l / mout * (f22 - f11) + f11
+        Mshift = -m_input / 2
+        self.Mshift = torch.exp(-1j * 2 * torch.pi * l * (Mshift + 0.5) / fs)
+
+    def transform(self, x, dim=-1):
+        x = x.to(self.device)
+        m = self.m_input
+
+        dim = dim if dim >= 0 else x.ndim + dim
+
+        if x.shape[dim] != m:
+            print(m)
+            print(x.shape)
+            raise ValueError(
+                f"Expected dimension {dim} to be of size {m}, but got {x.shape[dim]}"
+            )
+
+        x = x.transpose(dim, -1)
+
+        b_phase = self.b_phase.view((1,) * (x.ndim - 1) + (-1,))
+        x_weighted = x * b_phase
+
+        original_shape = x_weighted.shape
+        x_weighted = x_weighted.reshape(-1, m)
+
+        b_padded = torch.zeros(
+            (x_weighted.shape[0], self.padded_len),
+            dtype=torch.complex64,
+            device=self.device,
+        )
+        b_padded[:, :m] = x_weighted
+
+        b_fft = torch.fft.fft(b_padded, dim=1)
+        conv = b_fft * self.ft[None, :]
+        result = torch.fft.ifft(conv, dim=1)
+
+        result = (
+            result[:, self.m_input - 1 : self.mp] * self.h[self.m_input - 1 : self.mp]
+        )
+        result = result * self.Mshift[None, :]
+
+        new_shape = list(original_shape[:-1]) + [self.mout]
+        result = result.reshape(*new_shape)
+
+        result = result.transpose(-1, dim)
+
+        return result
+
+
+class DebyeWolf:
+    def __init__(
+        self,
+        Min,
+        xrange,
+        yrange,
+        zrange,
+        Mout,
+        lams,  # list of wavelengths
+        NA,
+        focal_length,
+        n1=1,
+        device="cpu",
+    ):
+        self.device = device
+        self.Min = Min
+        self.xrange = xrange
+        self.yrange = yrange
+        self.z_arr = torch.linspace(zrange[0], zrange[1], Mout[2], device=device)
+        self.Moutx, self.Mouty = Mout[0], Mout[1]
+        lams = torch.tensor(lams, device=device)
+        self.lams, self.k0, self.n1, self.NA, self.focal_length = (
+            lams,
+            2 * torch.pi / lams,
+            n1,
+            NA,
+            focal_length,
+        )
+
+        self.N = (Min - 1) / 2
+
+        m = torch.linspace(-Min / 2, Min / 2, Min, device=self.device)
+        n = torch.linspace(-Min / 2, Min / 2, Min, device=self.device)
+        self.m_grid, self.n_grid = torch.meshgrid(m, n, indexing="ij")
+
+        self.th = torch.asin(
+            torch.clamp(
+                NA * torch.sqrt(self.m_grid**2 + self.n_grid**2) / (self.N * n1), max=1
+            )
+        )
+        self.mask = self.th > torch.arcsin(torch.tensor(NA / n1))
+        self.phi = torch.atan2(self.n_grid, self.m_grid)
+        self.phi[self.phi < 0] += 2 * torch.pi
+
+        self._sqrt_costh = 1 / torch.sqrt(torch.cos(self.th).unsqueeze(-1))
+        self._sqrt_costh[torch.isnan(self._sqrt_costh)] = 0
+        self._sqrt_costh[self.mask] = 0
+
+        fs = lams * (Min - 1) / (2 * NA)
+        self.fs = fs
+        self.bluesteins_y = []
+        self.bluesteins_x = []
+        self.C = (
+            -1j
+            * torch.exp(1j * self.k0 * n1 * focal_length)
+            * focal_length
+            * (lams)
+            / (self.n1)
+            / fs
+            / fs
+        )
+        fs = fs.cpu().tolist()
+        for f in fs:
+            self.bluesteins_y.append(
+                BluesteinDFT(
+                    f / 2 + self.yrange[0],
+                    f / 2 + self.yrange[1],
+                    f,
+                    self.Mouty,
+                    Min,
+                    device=device,
+                )
+            )
+            self.bluesteins_x.append(
+                BluesteinDFT(
+                    f / 2 + self.xrange[0],
+                    f / 2 + self.xrange[1],
+                    f,
+                    self.Moutx,
+                    Min,
+                    device=device,
+                )
+            )
+        self.E_ideals = torch.ones_like(self.lams)
+        self.R = torch.stack(
+            [
+                -torch.sin(self.th) * torch.cos(self.phi),
+                -torch.sin(self.th) * torch.sin(self.phi),
+                torch.cos(self.th),
+            ],
+            dim=-1,
+        )
+        self.R = self.R.unsqueeze(-2)
+
+    def __call__(self, E, correct=False):
+        # The input E has shape (batch, x, y, 2, lam),
+        # where the z-component is not included.
+        # The output E has shape (batch, x, y, z, 3, lam).
+        # For different wavelengths (lam), a simple for-loop is used for now.
+
+        Ex_in, Ey_in = E[..., 0:1, :], E[..., 1:2, :]
+        th = self.th.unsqueeze(-1).unsqueeze(-1)
+        phi = self.phi.unsqueeze(-1).unsqueeze(-1)
+        z_arr = self.z_arr.unsqueeze(0).unsqueeze(0).unsqueeze(-1)
+        k0, n1 = self.k0.view(1, 1, 1, -1), self.n1
+
+        costh = torch.cos(th)
+        _sqrt_costh = self._sqrt_costh.unsqueeze(-1)
+        phase = torch.exp(1j * k0 * n1 * z_arr * costh)
+        deltadim = 0
+        C = (self.C / self.E_ideals).view(1, 1, 1, -1)
+        E_out = torch.zeros(
+            [self.Moutx, self.Mouty, self.z_arr.numel(), 3, self.lams.numel()],
+            dtype=torch.complex64,
+            device=self.device,
+        )
+        if E.dim() == 5:
+            th = th.unsqueeze(0)
+            phi = phi.unsqueeze(0)
+            z_arr = z_arr.unsqueeze(0)
+            k0 = k0.unsqueeze(0)
+            costh = costh.unsqueeze(0)
+            _sqrt_costh = _sqrt_costh.unsqueeze(0)
+            phase = phase.unsqueeze(0)
+            C = C.unsqueeze(0)
+            deltadim = 1
+            E_out = torch.zeros(
+                [
+                    E.size(0),
+                    self.Moutx,
+                    self.Mouty,
+                    self.z_arr.numel(),
+                    3,
+                    self.lams.numel(),
+                ],
+                dtype=torch.complex64,
+                device=self.device,
+            )
+        Ex = (
+            (
+                Ex_in * (1 + (costh - 1) * torch.cos(phi) ** 2)
+                + Ey_in * (costh - 1) * torch.cos(phi) * torch.sin(phi)
+            )
+            * phase
+            * _sqrt_costh
+        )
+
+        Ey = (
+            (
+                Ex_in * (costh - 1) * torch.cos(phi) * torch.sin(phi)
+                + Ey_in * (1 + (costh - 1) * torch.sin(phi) ** 2)
+            )
+            * phase
+            * _sqrt_costh
+        )
+
+        Ez = (
+            (Ex_in * torch.cos(phi) + Ey_in * torch.sin(phi))
+            * torch.sin(th)
+            * phase
+            * _sqrt_costh
+        )
+        if correct:
+            temp = torch.stack([Ex, Ey, Ez], dim=-1)
+            temp = temp - 0.5 * self.R * torch.sum(self.R * temp, dim=-1, keepdim=True)
+            Ex, Ey, Ez = temp[:, :, :, 0], temp[:, :, :, 1], temp[:, :, :, 2]
+        for i in range(self.lams.numel()):
+            E_out[..., 0, i] = self.bluesteins_x[i].transform(
+                self.bluesteins_y[i].transform(Ex[..., i], dim=1 + deltadim),
+                dim=0 + deltadim,
+            )
+            E_out[..., 1, i] = self.bluesteins_x[i].transform(
+                self.bluesteins_y[i].transform(Ey[..., i], dim=1 + deltadim),
+                dim=0 + deltadim,
+            )
+            E_out[..., 2, i] = self.bluesteins_x[i].transform(
+                self.bluesteins_y[i].transform(Ez[..., i], dim=1 + deltadim),
+                dim=0 + deltadim,
+            )
+
+        return C * E_out
+
+    def Get_Z_offset_Phase(self, z):
+        k0, n1 = self.k0, self.n1
+        costh = torch.cos(self.th)
+        phase = k0.view(1, 1, -1) * n1 * z * (-costh).unsqueeze(-1)
+        return phase
+
+
+def fft_circular_conv2d(E, G):
+    E_fft = torch.fft.fft2(E.permute(2, 0, 1))
+    G_fft = torch.fft.fft2(G.permute(2, 0, 1))
+    C_fft = E_fft * G_fft
+    C_ifft = torch.fft.ifft2(C_fft).permute(1, 2, 0)
+    return C_ifft
+
+
+def fourier_upsample2d(E: torch.Tensor, up_sample=1) -> torch.Tensor:
+    x, y = E.shape[0], E.shape[1]
+    Xn, Yn = x * up_sample, y * up_sample
+
+    # 2D FFT（仅前两维），正交归一避免缩放差异
+    F = torch.fft.fft2(E, dim=(0, 1), norm="ortho")
+    Fc = torch.fft.fftshift(F, dim=(0, 1))
+
+    # 频域补零到更大网格
+    Fp = torch.zeros((Xn, Yn) + E.shape[2:], dtype=E.dtype, device=E.device)
+
+    x0 = (Xn - x + 1) // 2
+    y0 = (Yn - y + 1) // 2
+    Fp[x0 : x0 + x, y0 : y0 + y, ...] = Fc
+
+    Fp = torch.fft.ifftshift(Fp, dim=(0, 1))
+    up_E_full = torch.fft.ifft2(Fp, dim=(0, 1), norm="ortho")
+
+    X_keep = (x - 1) * up_sample + 1
+    Y_keep = (y - 1) * up_sample + 1
+    up_E = up_E_full[:X_keep, :Y_keep, ...] * up_sample
+    return up_E
+
+
+def upsample_coords_xy(x: torch.Tensor, y: torch.Tensor, up_sample: int):
+    """
+    输入:
+      x: 原 x 轴坐标，形状 (Nx,)
+      y: 原 y 轴坐标，形状 (Ny,)
+      up_sample: 正整数，上采样倍数
+
+    返回:
+      x_up: 形状 ((Nx-1)*up + 1,)
+      y_up: 形状 ((Ny-1)*up + 1,)
+      （范围不变，步长缩小 1/up）
+    说明:
+      假设 x、y 为等间距采样（FFT/零填充本身也要求等间距）
+    """
+    if up_sample < 1:
+        raise ValueError("up_sample 必须 >= 1")
+    if up_sample == 1:
+        return x, y
+
+    if x.numel() < 2 or y.numel() < 2:
+        raise ValueError("x 和 y 至少需要 2 个点用于定义范围。")
+
+    # 直接按起点和终点线性插值（等间距）
+    dx = (x[1] - x[0]) / up_sample * 1.5 * (up_sample - 1)
+    dy = (y[1] - y[0]) / up_sample * 1.5 * (up_sample - 1)
+    x_up = torch.linspace(
+        x[0] + dx,
+        x[-1] + dx,
+        (x.numel() - 1) * up_sample + 1,
+        device=x.device,
+        dtype=x.dtype,
+    )
+    y_up = torch.linspace(
+        y[0] + dy,
+        y[-1] + dy,
+        (y.numel() - 1) * up_sample + 1,
+        device=y.device,
+        dtype=y.dtype,
+    )
+    return x_up, y_up
+
+
+class Luneburg:
+    # The calculation method for Ez is different and omitted here.
+    # The sampling theorem must be satisfied: T < lambda / 2
+    def __init__(
+        self,
+        uv_len,  # Integration region
+        x_range,
+        y_range,
+        z_range,
+        sampling_interval,  # Sampling interval, a 1D array; strictly followed in xy-directions,
+        # while z-direction is adaptively adjusted
+        lams,
+        focal_length,  # Evaluation is based on this focal length
+        n1=1,
+        up_sample=1,  # Integer > 1,
+        device="cuda",
+    ):
+        self.device = device
+        self.up_sample = up_sample
+        lams = torch.tensor(lams, device=device)
+        self.lams = lams
+        self.k = 2 * torch.pi * n1 / self.lams
+        self.n1 = n1
+        self.focal_length = focal_length
+        uv_len = uv_len - sampling_interval[0]
+        ux_arr = torch.arange(
+            -uv_len / 2 + x_range[0] - sampling_interval[0],
+            x_range[1] + uv_len / 2,
+            step=sampling_interval[0],
+            device=device,
+        )
+
+        vy_arr = torch.arange(
+            -uv_len / 2 + y_range[0] - sampling_interval[1],
+            y_range[1] + uv_len / 2,
+            step=sampling_interval[1],
+            device=device,
+        )
+        z_arr = torch.linspace(
+            z_range[0] + focal_length,
+            z_range[1] + focal_length,
+            int((z_range[1] - z_range[0]) / sampling_interval[2] / 2) * 2 + 1,
+            device=device,
+        )
+
+        center = (
+            round((x_range[0] + x_range[1]) / 2 / sampling_interval[0])
+            * sampling_interval[0]
+        )
+        u_arr = ux_arr - center
+        u_start_id = torch.argmin(torch.abs(u_arr + uv_len / 2))
+        u_end_id = torch.argmin(torch.abs(u_arr - uv_len / 2))
+        u_arr = u_arr[u_start_id : u_end_id + 1] + sampling_interval[0] / 2
+        du = torch.mean(u_arr)
+        u_arr = u_arr - du
+        center = (
+            round((y_range[0] + y_range[1]) / 2 / sampling_interval[1])
+            * sampling_interval[1]
+        )
+        v_arr = vy_arr - center
+        v_start_id = torch.argmin(torch.abs(v_arr + uv_len / 2))
+        v_end_id = torch.argmin(torch.abs(v_arr - uv_len / 2))
+        v_arr = v_arr[v_start_id : v_end_id + 1] + sampling_interval[1] / 2
+        dv = torch.mean(v_arr)
+        v_arr = v_arr - dv
+
+        x_start_id = torch.argmin(torch.abs(ux_arr - x_range[0]))
+        x_end_id = torch.argmin(torch.abs(ux_arr - x_range[1]))
+        y_start_id = torch.argmin(torch.abs(vy_arr - y_range[0]))
+        y_end_id = torch.argmin(torch.abs(vy_arr - y_range[1]))
+
+        x_arr = ux_arr[x_start_id : x_end_id + 1] - du
+        y_arr = vy_arr[y_start_id : y_end_id + 1] - dv
+
+        co_public = (
+            -1j
+            * sampling_interval[0]
+            * sampling_interval[1]
+            * n1
+            / self.lams.view(1, 1, 1, -1)
+        )  # / 4 / torch.pi
+        R_q = torch.sqrt(
+            (ux_arr.view(-1, 1, 1, 1)) ** 2
+            + (vy_arr.view(1, -1, 1, 1)) ** 2
+            + z_arr.view(1, 1, -1, 1) ** 2
+        )
+        G_2D = self.k * R_q * 1j
+        G_2D = (
+            torch.exp(G_2D)
+            / R_q**2
+            * co_public
+            * (1 - 1 / G_2D)
+            * torch.sinc(
+                n1
+                * ux_arr.view(-1, 1, 1, 1)
+                * sampling_interval[0]
+                / (R_q * lams.view(1, 1, 1, -1))
+            )
+            * torch.sinc(
+                n1
+                * vy_arr.view(1, -1, 1, 1)
+                * sampling_interval[1]
+                / (R_q * lams.view(1, 1, 1, -1))
+            )
+        )
+
+        self.G_2D = G_2D
+
+        self.u_arr = u_arr
+        self.v_arr = v_arr
+        self.x_arr, self.y_arr = upsample_coords_xy(x_arr, y_arr, up_sample)
+        self.z_arr = z_arr
+
+        self.x_num = x_end_id - x_start_id + 1
+        self.y_num = y_end_id - y_start_id + 1
+
+        self.u_pad_num = ux_arr.numel() - u_arr.numel()
+        self.v_pad_num = vy_arr.numel() - v_arr.numel()
+        self.dS = sampling_interval[0] * sampling_interval[1]
+
+    def __call__(self, E):
+        # E should have shape (u, v, lam)
+        if E.dim() == 3:
+            E_padded = E
+            pad = (
+                0,
+                0,  # No padding applied to the 2nd dimension (D2)
+                self.v_pad_num,
+                0,  # Extend on both sides of the 1st dimension (D1)
+                self.u_pad_num,
+                0,  # Extend on both sides of the 0th dimension (D0)
+            )
+
+            E_padded = F.pad(E_padded, pad, mode="constant", value=0)
+            E_out = torch.zeros(
+                [
+                    self.x_num,
+                    self.y_num,
+                    self.z_arr.numel(),
+                    self.lams.numel(),
+                ],
+                device=self.device,
+                dtype=torch.complex64,
+            )
+            for z in range(self.z_arr.numel()):
+                E_out[:, :, z, :] = fft_circular_conv2d(
+                    E_padded, self.G_2D[:, :, z] * self.z_arr[z]
+                )[
+                    : self.x_num,
+                    : self.y_num,
+                ]
+            E_out = fourier_upsample2d(E_out, self.up_sample)
+            return E_out
+
+    def Get_G_2D(self, x_f, y_f, z_f):  # Typically used for optimization and validation
+        z = self.focal_length + z_f  # Detection coordinate
+        co_public = -1j * self.dS * self.n1 / self.lams  # / 4 / torch.pi
+        R_q = torch.sqrt(
+            z**2
+            + (x_f - self.u_arr.view(-1, 1, 1)) ** 2
+            + (y_f - self.v_arr.view(1, -1, 1)) ** 2
+        )
+        ik0R = self.k.view(1, 1, -1) * R_q * 1j
+        G_2D = torch.exp(ik0R) / R_q**2 * (1 - 1 / ik0R) * co_public
+
+        return z * G_2D
+
+
+class LuneburgOld:
+    # The calculation method for Ez is different and omitted here.
+    # The sampling theorem must be satisfied: T < lambda / 2
+    def __init__(
+        self,
+        uv_len,  # Integration region
+        x_range,
+        y_range,
+        z_range,
+        sampling_interval,  # Sampling interval, a 1D array; strictly followed in xy-directions,
+        # while z-direction is adaptively adjusted
+        lams,
+        focal_length,  # Evaluation is based on this focal length
+        n1=1,
+        scale=1,  # Integer > 1, used to represent a reduced sampling frequency for E
+        device="cuda",
+    ):
+        self.device = device
+        self.scale = scale
+        lams = torch.tensor(lams, device=device)
+        self.lams = lams
+        self.k = 2 * torch.pi * n1 / self.lams
+        self.n1 = n1
+        self.focal_length = focal_length
+        uv_len = uv_len - sampling_interval[0]
+        ux_arr = torch.arange(
+            -uv_len / 2 + x_range[0] - sampling_interval[0],
+            x_range[1] + uv_len / 2,
+            step=sampling_interval[0],
+            device=device,
+        )
+
+        vy_arr = torch.arange(
+            -uv_len / 2 + y_range[0] - sampling_interval[1],
+            y_range[1] + uv_len / 2,
+            step=sampling_interval[1],
+            device=device,
+        )
+        z_arr = torch.linspace(
+            z_range[0] + focal_length,
+            z_range[1] + focal_length,
+            int((z_range[1] - z_range[0]) / sampling_interval[2] / 2) * 2 + 1,
+            device=device,
+        )
+
+        center = (
+            round((x_range[0] + x_range[1]) / 2 / sampling_interval[0])
+            * sampling_interval[0]
+        )
+        u_arr = ux_arr - center
+        u_start_id = torch.argmin(torch.abs(u_arr + uv_len / 2))
+        u_end_id = torch.argmin(torch.abs(u_arr - uv_len / 2))
+        u_arr = u_arr[u_start_id : u_end_id + 1 : scale] + sampling_interval[0] / 2
+        du = torch.mean(u_arr)
+        u_arr = u_arr - du
+        center = (
+            round((y_range[0] + y_range[1]) / 2 / sampling_interval[1])
+            * sampling_interval[1]
+        )
+        v_arr = vy_arr - center
+        v_start_id = torch.argmin(torch.abs(v_arr + uv_len / 2))
+        v_end_id = torch.argmin(torch.abs(v_arr - uv_len / 2))
+        v_arr = v_arr[v_start_id : v_end_id + 1 : scale] + sampling_interval[1] / 2
+        dv = torch.mean(v_arr)
+        v_arr = v_arr - dv
+
+        x_start_id = torch.argmin(torch.abs(ux_arr - x_range[0]))
+        x_end_id = torch.argmin(torch.abs(ux_arr - x_range[1]))
+        y_start_id = torch.argmin(torch.abs(vy_arr - y_range[0]))
+        y_end_id = torch.argmin(torch.abs(vy_arr - y_range[1]))
+
+        x_arr = ux_arr[x_start_id : x_end_id + 1] - du
+        y_arr = vy_arr[y_start_id : y_end_id + 1] - dv
+
+        co_public = (
+            -1j
+            * sampling_interval[0]
+            * sampling_interval[1]
+            * n1
+            / self.lams.view(1, 1, 1, -1)
+        )  # / 4 / torch.pi
+        R_q = torch.sqrt(
+            (ux_arr.view(-1, 1, 1, 1)) ** 2
+            + (vy_arr.view(1, -1, 1, 1)) ** 2
+            + z_arr.view(1, 1, -1, 1) ** 2
+        )
+        G_2D = self.k * R_q * 1j
+        G_2D = torch.exp(G_2D) / R_q**2 * co_public * (1 - 1 / G_2D)
+
+        self.G_2D = G_2D
+
+        self.u_arr = u_arr
+        self.v_arr = v_arr
+        self.x_arr = x_arr
+        self.y_arr = y_arr
+        self.z_arr = z_arr
+
+        self.x_num = x_end_id - x_start_id + 1
+        self.y_num = y_end_id - y_start_id + 1
+
+        self.u_pad_num = ux_arr.numel() - u_arr.numel() * scale
+        self.v_pad_num = vy_arr.numel() - v_arr.numel() * scale
+        self.dS = sampling_interval[0] * sampling_interval[1]
+
+    def __call__(self, E):
+        # E should have shape (u, v, lam)
+        if E.dim() == 3:
+            E_padded = E.repeat_interleave(self.scale, dim=0)
+            E_padded = E_padded.repeat_interleave(self.scale, dim=1)
+            pad = (
+                0,
+                0,  # No padding applied to the 2nd dimension (D2)
+                self.v_pad_num,
+                0,  # Extend on both sides of the 1st dimension (D1)
+                self.u_pad_num,
+                0,  # Extend on both sides of the 0th dimension (D0)
+            )
+
+            E_padded = F.pad(E_padded, pad, mode="constant", value=0)
+            E_out = torch.zeros(
+                [
+                    self.x_num,
+                    self.y_num,
+                    self.z_arr.numel(),
+                    self.lams.numel(),
+                ],
+                device=self.device,
+                dtype=torch.complex64,
+            )
+            for z in range(self.z_arr.numel()):
+                E_out[:, :, z, :] = fft_circular_conv2d(
+                    E_padded, self.G_2D[:, :, z] * self.z_arr[z]
+                )[
+                    : self.x_num,
+                    : self.y_num,
+                ]
+            return E_out
+
+    def Get_G_2D(self, x_f, y_f, z_f):  # Typically used for optimization and validation
+        z = self.focal_length + z_f  # Detection coordinate
+        co_public = (
+            -1j * self.dS * self.n1 / self.lams * self.scale**2
+        )  # / 4 / torch.pi
+        R_q = torch.sqrt(
+            z**2
+            + (x_f - self.u_arr.view(-1, 1, 1)) ** 2
+            + (y_f - self.v_arr.view(1, -1, 1)) ** 2
+        )
+        ik0R = self.k.view(1, 1, -1) * R_q * 1j
+        G_2D = torch.exp(ik0R) / R_q**2 * (1 - 1 / ik0R) * co_public
+
+        return z * G_2D

torchopticsy-0.8.3/torchopticsy/Interference.py

@@ -0,0 +1,284 @@
+# %%
+# 滤波确实可以去掉伪影
+import torch
+import torch.nn.functional as F
+import cv2
+import os
+import matplotlib.pyplot as plt
+
+# pixel_size=2.4/100*1.8/2
+# class OffAxisPhaseOld:
+#     def __init__(
+#         self,
+#         file_name,
+#         continuous=True,  # 默认是连续光
+#         block_size=1024,  # 每个子块的大小
+#         overlap=0,  # 子块之间是否重叠
+#     ):
+#         self.continuous = continuous
+#         self.block_size = block_size
+#         self.overlap = overlap
+
+#         # 支持的图片后缀
+#         self.valid_ext = [".jpg", ".png", ".bmp"]
+
+#         # 读取图像
+#         BACK = self._read_image("BACK", file_name)
+#         self.OBJ = torch.clamp(self._read_image("OBJ", file_name) - BACK, min=0)
+#         self.REF = torch.clamp(self._read_image("REF", file_name) - BACK, min=0)
+#         self.OBJ_REF = torch.clamp(self._read_image("OBJ_REF", file_name) - BACK, min=0)
+#         self.INC = torch.clamp(self._read_image("INC", file_name) - BACK, min=0)
+#         self.INC_REF = torch.clamp(self._read_image("INC_REF", file_name) - BACK, min=0)
+
+#     def _read_image(self, prefix, file_name):
+#         for ext in self.valid_ext:
+#             img_path = os.path.join(file_name, f"{prefix}{ext}")
+#             if os.path.exists(img_path):
+#                 img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
+#                 if img is None:
+#                     raise ValueError(f"无法读取图像: {img_path}")
+#                 img_tensor = torch.from_numpy(img)
+
+#                 if img_tensor.dtype == torch.uint8:
+#                     img_tensor = img_tensor.to(torch.float32) / 255.0
+#                 elif img_tensor.dtype == torch.uint16:
+#                     img_tensor = img_tensor.to(torch.float32) / 65535.0
+#                 else:
+#                     raise ValueError(f"不支持的图像位深: {img_tensor.dtype}")
+#                 return img_tensor
+#         raise FileNotFoundError(f"在 {file_name} 中找不到 {prefix} 图像")
+
+#     def _split_blocks(self, img):
+#         H, W = img.shape
+#         step = self.block_size - self.overlap
+#         blocks = []
+
+#         y_list = list(range(0, H - self.block_size + 1, step))
+#         x_list = list(range(0, W - self.block_size + 1, step))
+
+#         # 确保边缘覆盖
+#         if y_list[-1] != H - self.block_size:
+#             y_list.append(H - self.block_size)
+#         if x_list[-1] != W - self.block_size:
+#             x_list.append(W - self.block_size)
+
+#         for y in y_list:
+#             for x in x_list:
+#                 blocks.append(
+#                     ((y, x), img[y : y + self.block_size, x : x + self.block_size])
+#                 )
+
+#         return blocks, H, W
+
+#     def _estimate_angle(self, block, global_angle):
+#         """先用全局角度粗筛，再精确估计局部峰值角度"""
+#         f = torch.fft.fftshift(torch.fft.fft2(block))
+
+#         H, W = f.shape
+#         ky_vals = torch.fft.fftshift(torch.fft.fftfreq(H, d=pixel_size)).to(f.device)
+#         kx_vals = torch.fft.fftshift(torch.fft.fftfreq(W, d=pixel_size)).to(f.device)
+#         KY, KX = torch.meshgrid(ky_vals, kx_vals, indexing="ij")
+
+#         angle = torch.deg2rad(torch.tensor(global_angle))
+#         vx, vy = torch.cos(angle), torch.sin(angle)
+
+#         cross = vx * KY - vy * KX
+#         mask = cross >= 0
+#         mag = torch.abs(f) * mask  # 只保留一半
+#         # 去掉直流分量
+#         cy, cx = H // 2, W // 2
+#         mag[cy - 20 : cy + 20, cx - 20 : cx + 20] = 0
+#         # ==== Step 2: 找峰值 ====
+#         max_pos = torch.nonzero(mag == mag.max(), as_tuple=False)[0]
+#         dy, dx = max_pos[0] - mag.size(0) / 2, max_pos[1] - mag.size(1) / 2
+#         angle_loc = torch.atan2(-dx, dy)  # 子图峰值方向
+#         Rmax = torch.sqrt(dy**2+dx**2)*0.75#滤波半径
+#         kmax=torch.sqrt(ky_vals[max_pos[0]]**2+kx_vals[max_pos[1]]**2)
+
+#         print(kmax)
+#         return torch.rad2deg(angle_loc).item(), f, mag,Rmax
+
+#     def _filter_one(self, input_tensor, angle_deg):
+#         f = torch.fft.fftshift(torch.fft.fft2(input_tensor))
+
+#         H, W = f.shape
+#         ky_vals = torch.fft.fftshift(torch.fft.fftfreq(H, d=1.0 / H)).to(f.device)
+#         kx_vals = torch.fft.fftshift(torch.fft.fftfreq(W, d=1.0 / W)).to(f.device)
+#         KY, KX = torch.meshgrid(ky_vals, kx_vals, indexing="ij")
+
+#         angle = torch.deg2rad(torch.tensor(angle_deg))
+#         vx, vy = torch.cos(angle), torch.sin(angle)
+
+#         cross = vx * KY - vy * KX
+#         mask = cross >= 0
+
+#         f_filtered = f * mask
+#         result = torch.fft.ifft2(torch.fft.ifftshift(f_filtered))
+#         return result, f, mask
+
+#     def __call__(self, angle_deg=0, visualize=False, vis_num=1,lowpass=False):
+#         # 分块处理
+#         a_full = self.OBJ_REF - self.OBJ - self.REF
+#         b_full = self.INC_REF - self.INC - self.REF
+
+#         blocks, H, W = self._split_blocks(a_full)
+#         E_full = torch.zeros((H, W), dtype=torch.complex64)
+
+#         vis_count = 0
+#         x=torch.arange(self.block_size).float()
+#         x-=torch.mean(x)
+#         y=torch.arange(self.block_size).float()
+#         y-=torch.mean(y)
+#         R=torch.sqrt(x.view(-1,1)**2+y.view(1,-1)**2)
+
+#         for (y, x), a_block in blocks:
+#             b_block = b_full[y : y + self.block_size, x : x + self.block_size]
+
+#             # 每个子块独立估计方向
+#             block_angle, f_b, mag_b,Rmax = self._estimate_angle(b_block, angle_deg)
+
+#             a_filtered, f_a, mask_a = self._filter_one(a_block, block_angle)
+#             b_filtered, _, mask_b = self._filter_one(b_block, block_angle)
+
+#             E = a_filtered/ (b_filtered)
+#             E[~torch.isfinite(E)] = 0
+
+#             # 放回全图
+#             if lowpass:
+#                 E = torch.fft.fftshift(torch.fft.fft2(E))
+#                 plt.figure(figsize=(8,8))
+#                 plt.subplot(2, 2, 1)
+#                 plt.pcolormesh(torch.log1p(torch.abs(E[800:-800,800:-800])).cpu())
+#                 plt.colorbar()
+#                 plt.title("E FFT magnitude (before)")
+#                 E[R>Rmax]=0
+#                 plt.subplot(2, 2, 2)
+#                 plt.pcolormesh(torch.log1p(torch.abs(E[800:-800,800:-800])).cpu())
+#                 plt.colorbar()
+#                 plt.title("E FFT magnitude (after)")
+#                 plt.subplot(2, 2, 3)
+#                 plt.pcolormesh((R[800:-800,800:-800]>Rmax).float().cpu())
+#                 plt.title("R magnitude (after)")
+
+
+#                 E = torch.fft.ifft2(torch.fft.ifftshift(E))
+#             E_full[y : y + self.block_size, x : x + self.block_size] = E
+
+#             # 可视化部分子图
+#             if visualize and vis_count < vis_num:
+#                 plt.figure(figsize=(12, 6))
+#                 plt.suptitle(f"Block ({y},{x}), angle={block_angle:.2f}°")
+
+#                 plt.subplot(2, 2, 1)
+#                 plt.imshow(torch.log1p(mag_b).cpu(), cmap="gray")
+#                 plt.title("b FFT magnitude (before)")
+
+#                 plt.subplot(2, 2, 2)
+#                 plt.imshow(torch.log1p(torch.abs(f_b * mask_b)).cpu(), cmap="gray")
+#                 plt.title("b FFT after mask")
+
+#                 plt.subplot(2, 2, 3)
+#                 plt.imshow(torch.log1p(torch.abs(torch.abs(f_a))).cpu(), cmap="gray")
+#                 plt.title("a FFT magnitude (before)")
+
+#                 plt.subplot(2, 2, 4)
+#                 plt.imshow(torch.log1p(torch.abs(f_a * mask_a)).cpu(), cmap="gray")
+#                 plt.title("a FFT after mask")
+
+#                 plt.tight_layout()
+#                 plt.show()
+
+#                 vis_count += 1
+
+#         return E_full
+
+
+class OffAxisPhase:
+    def __init__(self, file_name, pixel_size=2.4 / 100 * 1.8 / 2):
+        self.pixel_size = pixel_size
+
+        # 支持的图片后缀
+        self.valid_ext = [".jpg", ".png", ".bmp"]
+
+        # 读取图像
+        BACK = self._read_image("BACK", file_name)
+        self.OBJ = torch.clamp(self._read_image("OBJ", file_name) - BACK, min=0)
+        self.REF = torch.clamp(self._read_image("REF", file_name) - BACK, min=0)
+        self.OBJ_REF = torch.clamp(self._read_image("OBJ_REF", file_name) - BACK, min=0)
+        self.INC = torch.clamp(self._read_image("INC", file_name) - BACK, min=0)
+        self.INC_REF = torch.clamp(self._read_image("INC_REF", file_name) - BACK, min=0)
+
+    def _read_image(self, prefix, file_name):
+        for ext in self.valid_ext:
+            img_path = os.path.join(file_name, f"{prefix}{ext}")
+            if os.path.exists(img_path):
+                img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
+                if img is None:
+                    raise ValueError(f"无法读取图像: {img_path}")
+                img_tensor = torch.from_numpy(img)
+
+                if img_tensor.dtype == torch.uint8:
+                    img_tensor = img_tensor.to(torch.float32) / 255.0
+                elif img_tensor.dtype == torch.uint16:
+                    img_tensor = img_tensor.to(torch.float32) / 65535.0
+                else:
+                    raise ValueError(f"不支持的图像位深: {img_tensor.dtype}")
+                return img_tensor
+        raise FileNotFoundError(f"在 {file_name} 中找不到 {prefix} 图像")
+
+    def _get_mask(self, block, global_angle):
+        """估计局部峰值"""
+        f = torch.fft.fftshift(torch.fft.fft2(block))
+        H, W = f.shape
+        ky_vals = torch.fft.fftshift(torch.fft.fftfreq(H, d=self.pixel_size)).to(
+            f.device
+        )
+        kx_vals = torch.fft.fftshift(torch.fft.fftfreq(W, d=self.pixel_size)).to(
+            f.device
+        )
+        KY, KX = torch.meshgrid(ky_vals, kx_vals, indexing="ij")
+
+        angle = torch.deg2rad(torch.tensor(global_angle))
+        vx, vy = torch.cos(angle), torch.sin(angle)
+
+        cross = vx * KY - vy * KX
+        mask = cross >= 0
+        mag = torch.abs(f) * mask  # 只保留一半
+        # 去掉直流分量
+        cy, cx = H // 2, W // 2
+        mag[cy - 15 : cy + 15, cx - 15 : cx + 15] = 0
+        # ==== Step 2: 找峰值 ====
+        max_pos = torch.nonzero(mag == mag.max(), as_tuple=False)[0]
+        kyc, kxc = ky_vals[max_pos[0]], kx_vals[max_pos[1]]
+        kmax = torch.sqrt(kyc**2 + kxc**2) * 0.9  # 滤波半径
+        print("kmax:" + str(round(kmax.cpu().item(), 2)))
+        mask = torch.sqrt((KY - kyc) ** 2 + (KX - kxc) ** 2) < kmax
+        return mask
+
+    def __call__(self, angle_deg=0, visualize=False):
+        # 分块处理
+        a_full = self.OBJ_REF - self.OBJ - self.REF
+        b_full = self.INC_REF - self.INC - self.REF
+
+        mask = self._get_mask(b_full, angle_deg)
+        a_full = torch.fft.fftshift(torch.fft.fft2(a_full))
+        a_full[~mask] = 0
+        a_full = torch.fft.ifft2(torch.fft.ifftshift(a_full))
+
+        b_full = torch.fft.fftshift(torch.fft.fft2(b_full))
+        if visualize:
+            plt.figure(figsize=(12, 6))
+            plt.subplot(2, 1, 1)
+            plt.imshow(torch.log1p(torch.abs(b_full)).cpu(), cmap="gray")
+            plt.title("b_full")
+
+            plt.subplot(2, 1, 2)
+            plt.imshow(mask.cpu(), cmap="gray")
+            plt.title("mask")
+        b_full[~mask] = 0
+        b_full = torch.fft.ifft2(torch.fft.ifftshift(b_full))
+
+        E = a_full / (b_full)
+        E[~torch.isfinite(E)] = 0
+
+        return E

torchopticsy-0.8.3/torchopticsy/__init__.py

@@ -0,0 +1 @@
+

torchopticsy-0.8.3/torchopticsy.egg-info/PKG-INFO

@@ -0,0 +1,26 @@
+Metadata-Version: 2.1
+Name: torchopticsy
+Version: 0.8.3
+Summary: PyTorch-based optics caculation
+Author: YuningYe
+Author-email: 1956860113@qq.com
+License: MIT
+Description-Content-Type: text/markdown
+Requires-Dist: torch
+Requires-Dist: opencv-python
+Requires-Dist: matplotlib
+
+# torchOpticsY
+A PyTorch-based optics calculation library.
+
+
+## Features
+- Luneburg integral 
+- Debye鈥揥olf integral.
+
+## Install
+pip install torchOpticsY
+
+
+0.2 - Interference: OffAxisPhase
+0.1 - Diffraction: Debye_Wolf,Luneburg

torchopticsy-0.8.3/torchopticsy.egg-info/SOURCES.txt

@@ -0,0 +1,9 @@
+setup.py
+torchopticsy/Diffraction.py
+torchopticsy/Interference.py
+torchopticsy/__init__.py
+torchopticsy.egg-info/PKG-INFO
+torchopticsy.egg-info/SOURCES.txt
+torchopticsy.egg-info/dependency_links.txt
+torchopticsy.egg-info/requires.txt
+torchopticsy.egg-info/top_level.txt

torchopticsy-0.8.3/torchopticsy.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

torchopticsy-0.8.3/torchopticsy.egg-info/requires.txt

@@ -0,0 +1,3 @@
+torch
+opencv-python
+matplotlib

torchopticsy-0.8.3/torchopticsy.egg-info/top_level.txt

@@ -0,0 +1 @@
+torchopticsy