aopera 0.1.0__py3-none-any.whl

aopera/ffwfs.py

@@ -0,0 +1,335 @@
+"""
+Created on Fri Nov 11 19:00:43 2022
+Fourier-Filtering Wavefront Sensor
+@author: v.chambouleyron r.fetick
+"""
+
+import numpy as np
+from scipy.signal import fftconvolve
+from numpy.fft import fft2, fftshift, rfft2, irfft2
+from aopera.readconfig import read_config_file, read_config_tiptop, set_attribute
+from functools import lru_cache
+import logging
+
+#%% FF-WFS BASIC FUNCTIONS
+
+@lru_cache(maxsize=2)
+def modulation(nx, rmod, samp):
+    """
+    Build the modulation ring
+    
+    Parameters
+    ----------
+    nx : number of pixels.
+    rmod : modulation radius [lambda/D].
+    samp : sampling (2=Shannon).
+    """
+    rmod_px = rmod*samp
+    xx,yy = np.mgrid[0:nx,0:nx] - nx//2
+    rr = np.sqrt(xx**2+yy**2)
+    return (rr<(rmod_px+0.5))*(rr>=(rmod_px-0.5))
+
+
+@lru_cache(maxsize=2)
+def pyramid_mask(nx, nFaces, angle=0.35, theta_ref=0, four_pixels=False):
+    """
+    Build the pyramid mask
+    
+    Parameters
+    ----------
+    nx : number of pixels.
+    nFaces : number of faces for the pyramidal mask.
+    angle : angle of the pyramid  #TODO: definition? which unit?
+    theta_ref : rotation angle of the pyramid faces.
+    """
+    if four_pixels:
+        center = nx//2 - 0.5
+    else:
+        center = nx//2
+    xx,yy = np.mgrid[0:nx,0:nx] - center
+    theta = np.mod(np.arctan2(xx,yy)+np.pi-theta_ref, 2*np.pi)
+    theta_face = 2*np.pi/nFaces
+    msk = np.zeros((nx,nx), dtype=complex)
+
+    for k in range(0,nFaces):
+        # Tesselation
+        theta_bnd = ((k*theta_face)<=theta) * (theta<((k+1)*theta_face))
+        # Tip-Tilt 
+        theta_direction = (k+0.5)*theta_face
+        c_tip = np.sin(theta_direction+theta_ref)
+        c_tilt = np.cos(theta_direction+theta_ref)
+        # Complex mask on each face
+        msk += theta_bnd * np.exp(2j*np.pi*angle*(c_tip*xx+c_tilt*yy))
+        
+    return msk
+    
+
+def ffwfs_impulse_response(mask, modu, psf):
+    """
+    Compute the impulse response of a FF-WFS.
+    """
+    wtot = fftshift(irfft2(rfft2(modu)*rfft2(psf), s=psf.shape))
+    wtot = wtot/np.sum(wtot)
+    mask_ft = fft2(fftshift(mask)) / np.size(mask)
+    return 2*np.imag(np.conj(mask_ft)*fft2(fftshift(mask*wtot)))
+
+
+def ffwfs_transfer_function(*args, **kwargs):
+    """
+    Compute the transfer function of a FF-WFS.
+    """
+    return fftshift(fft2(ffwfs_impulse_response(*args, **kwargs)))
+
+
+def ffwfs_sensitivity(mask, modu, psf_calib, psf):
+    """
+    Compute the sensitivity map of a FF-WFS in the (fx,fy) domain.
+    
+    Parameters
+    ----------
+    mask : the complex mask of the FF-WFS.
+    modu : the modulation circle mask.
+    psf_calib : PSF used as reference for calibration.
+    psf : current PSF shape to compute optical gains.
+    """
+    otf_calib = rfft2(psf_calib)
+    TF = ffwfs_transfer_function(mask, modu, psf)
+    return np.sqrt(np.real(fftshift(irfft2(rfft2(np.abs(TF)**2)*otf_calib, s=psf.shape))))
+
+
+def ffwfs_reconstructor(*args, thresh=1e-2):
+    """
+    Compute the reconstructor map of a FF-WFS in the (fx,fy) domain.
+    See also: ffwfs_sensitivity
+    """
+    sensi = ffwfs_sensitivity(*args)
+    rec = np.zeros(sensi.shape)
+    vld = np.where(sensi>thresh)
+    rec[vld] = 1/sensi[vld]
+    return rec
+
+
+def ffwfs_visibility(*args, **kwargs):
+    """
+    Compute the FF-WFS visible frequencies.
+    See also: ffwfs_sensitivity, ffwfs_reconstructor
+    """
+    sensi = ffwfs_sensitivity(*args)
+    rec = ffwfs_reconstructor(*args, **kwargs)
+    return np.real(sensi*rec)
+
+
+def optical_gain(mask, modu, psf_calib, psf_onsky, obj_sky=None, modu_sky=True):
+    """
+    Compute optical gains map in the (fx,fy) domain.
+    
+    Parameters
+    ----------
+    mask : the complex mask of the FF-WFS.
+    modu : the modulation circle mask.
+    psf_calib : PSF used as reference for calibration.
+    psf_onsky : current PSF shape to compute optical gains.
+    
+    Keywords
+    --------
+    obj_sky : the extended object image (if applicable)
+    modu_sky : should we modulate onsky?
+    """
+    if modu_sky:
+        modu_onsky = modu
+    else:
+        modu_onsky = modulation(modu.shape[0], 0, 2) # Dirac
+    if obj_sky is not None:
+        psf_onsky = fftconvolve(psf_onsky, obj_sky/np.sum(obj_sky), mode='same')
+    sensi_calib = ffwfs_sensitivity(mask, modu, psf_calib, psf_calib)
+    sensi_sky   = ffwfs_sensitivity(mask, modu_onsky, psf_calib, psf_onsky)
+    return sensi_sky/sensi_calib
+
+
+#%% POWER SPECTRAL DENSITIES
+
+def ffwfs_psd_aliasing(fx, fy, psd_input, ffwfsrec, pix):
+    """
+    Aliasing PSD of a FF-WFS.
+    Assume a constant FF-WFS sensitivity at high-frequency.
+    
+    Parameters
+    ----------
+    fx : np.array
+        The spatial frequencies on X [1/m].
+    fy : np.array
+        The spatial frequencies on Y [1/m].
+    psd_input : function
+        The input PSD to be aliased. Should be called as: psd=psd_input(fx,fy)
+    ffwfsrec : np.array
+        see ffwfs_reconstructor
+    pix : float
+        Pixel size in the pupil plane [m].
+    """
+    psd_alias = np.zeros(fx.shape)
+    fcut = 1/(2*pix)
+    nshift = 3
+    ffwfssensi = 1/np.mean(ffwfsrec[0,:]) # constant sensitivity at HF
+    psdfilter = np.abs(ffwfsrec*ffwfssensi)**2
+    for i in range(-nshift,nshift+1):
+        for j in range(-nshift,nshift+1):
+            if (i!=0) or (j!=0):
+                fx_shift = fx - i*2*fcut
+                fy_shift = fy - j*2*fcut
+                psd_atmo_shift = psd_input(fx_shift,fy_shift)
+                pix_filter = np.sinc(fx_shift*pix)*np.sinc(fy_shift*pix)
+                psd_alias += psd_atmo_shift * psdfilter * pix_filter**2
+    return psd_alias
+
+
+def psd_pyrwfs_noise_ron(mask, modu, nphot, var_ron, nssp, psf_onsky, psf_calib):
+    """
+    Compute the phase PSD due to the AO RON.
+    The PSD is given on the (fx,fy) after reconstruction of the phase,
+    but before the loop temporal noise filtering.
+    
+    Parameters
+    ----------
+    mask : the complex mask of the Pyramid-WFS.
+    modu : the modulation circle mask.
+    nphot : float (>=0)
+        Number of photons available for the neasurement in total.
+    var_ron : camera RON in photons unit/frame/px
+    nssp = number of pixels in one pupil image
+    psf_onsky : current PSF shape to compute sensitivity on-sky.
+    psf_calib : PSF used as reference for calibration.
+    
+    Return
+    ------
+    the AO error PSD [rad²m²].
+    The PSD is given at the measurement wvl !!!
+    To get it at the scientific wvl, it should be multiplied by (wvl_wfs/wvl_sci)**2.
+    
+    References
+    ----------
+    V. Chambouleyron, 2021, PhD thesis: section 3.1
+    """
+    sensi_ron = ffwfs_sensitivity(mask,modu,psf_calib,psf_onsky)
+    var_ron = (nssp*var_ron)/nphot**2 * (1/sensi_ron)**2
+    return var_ron
+
+
+def psd_pyrwfs_noise_photon(mask, nfaces, modu, nphot, psf_onsky, psf_calib, emccd=False):
+    """
+    Compute the phase PSD due to the photon noise.
+    The PSD is given on the (fx,fy) after reconstruction of the phase,
+    but before the loop temporal noise filtering.
+    
+    Parameters
+    ----------
+    mask : the complex mask of the Pyramid-WFS.
+    nfaces : number of faces of the Pyramid-WFS
+    modu : the modulation circle mask.
+    nphot : float (>=0)
+        Number of photons available for the neasurement in total.
+    psf_onsky : current PSF shape to compute sensitivity on-sky.
+    
+    Keywords
+    --------
+    emccd : boolean
+        Multiply result by excess factor if camera is EMCCD
+    
+    Return
+    ------
+    the AO error PSD [rad²m²].
+    The PSD is given at the measurement wvl !!!
+    To get it at the scientific wvl, it should be multiplied by (wvl_wfs/wvl_sci)**2.
+    
+    References
+    ----------
+    V. Chambouleyron, 2021, PhD thesis: section 3.1
+    """
+    #FIXME: this sensitivity gives correct results wrt OOPAO, but why?
+    logging.debug('FFWFS sensitivity has been increased to match OOPAO.')
+    sensi_phot = ffwfs_sensitivity(mask,modu,psf_calib,psf_onsky) / np.sqrt(nfaces/2.5) #/ np.sqrt(nfaces) # Rough formula, to be checked more !!!
+    var_phot = 1/nphot*(1/sensi_phot)**2
+    if emccd:
+        var_phot *= 2 # excess factor squared
+    return var_phot
+
+#%% FFWFS CLASS
+class PWFS:
+    def __init__(self, dictionary):
+        self.nface = 4
+        self.angle = 0.35
+        set_attribute(self, dictionary, 'pwfs')
+           
+    def __repr__(self):
+        s  = 'aopera.PWFS\n'
+        s += '------------\n'
+        s += 'lenslet : %u \n'%self.lenslet
+        s += 'ron     : %.1f e-\n'%self.ron
+        s += 'EMCCD   : %s\n'%self.emccd
+        s += 'R_modul : %.1f l/D\n'%self.modulation
+        s += 'OG comp : %s\n'%self.og_compensation
+        return s
+           
+    @staticmethod
+    def from_file(filepath, category='pwfs'):
+        return PWFS(read_config_file(filepath, category))
+    
+    @staticmethod
+    def from_file_tiptop(filepath):
+        return PWFS(read_config_tiptop(filepath)[2])
+    
+    @staticmethod
+    def from_oopao(wfs):
+        return PWFS({'lenslet':wfs.nSubap,
+                     'ron':wfs.cam.readoutNoise,
+                     'emccd':wfs.cam.sensor=='EMCCD',
+                     'modulation':wfs.modulation})
+    
+    def modulation_mask(self, nx, samp):
+        return modulation(nx, self.modulation, samp)
+    
+    def pyramid_mask(self, nx, theta_ref=0):
+        return pyramid_mask(nx, self.nface, angle=self.angle, theta_ref=theta_ref)
+    
+    def sensitivity(self, samp, psf_calib, psf):
+        nx = psf.shape[0]
+        mask = self.pyramid_mask(nx)
+        modu = self.modulation_mask(nx, samp)
+        return ffwfs_sensitivity(mask, modu, psf_calib, psf)
+    
+    def reconstructor(self, samp, psf_calib, thresh=1e-2):
+        nx = psf_calib.shape[0]
+        mask = self.pyramid_mask(nx)
+        modu = self.modulation_mask(nx, samp)
+        return ffwfs_reconstructor(mask, modu, psf_calib, psf_calib, thresh=thresh)
+    
+    def visibility(self, samp, psf_calib, thresh=1e-2):
+        nx = psf_calib.shape[0]
+        mask = self.pyramid_mask(nx)
+        modu = self.modulation_mask(nx, samp)
+        return ffwfs_visibility(mask, modu, psf_calib, psf_calib, thresh=thresh)
+    
+    def optical_gain(self, samp, psf_calib, psf_onsky, **kwargs):
+        nx = psf_calib.shape[0]
+        mask = self.pyramid_mask(nx)
+        modu = self.modulation_mask(nx, samp)
+        return optical_gain(mask, modu, psf_calib, psf_onsky, **kwargs)
+    
+    def psd_aliasing(self, fx, fy, psd_input, ffwfsrec, diameter, nact=np.inf):
+        # NOTE : small subtelty on <nact> :
+        # if a mode is not shown to the WFS by the DM, it cannot
+        # be disantangled by the reconstructor, so it is aliased.
+        return ffwfs_psd_aliasing(fx, fy, psd_input, ffwfsrec, diameter/min(self.lenslet,nact-1))
+    
+    def psd_noise_ron(self, samp, nphot, psf_onsky, psf_calib):
+        nx = psf_calib.shape[0]
+        mask = self.pyramid_mask(nx)
+        modu = self.modulation_mask(nx, samp)
+        nssp = np.pi*(self.lenslet/2)**2
+        return psd_pyrwfs_noise_ron(mask, modu, nphot, self.ron**2, nssp, psf_onsky, psf_calib)
+    
+    def psd_noise_photon(self, samp, nphot, psf_onsky, psf_calib):
+        nx = psf_calib.shape[0]
+        mask = self.pyramid_mask(nx)
+        modu = self.modulation_mask(nx, samp)
+        return psd_pyrwfs_noise_photon(mask, self.nface, modu, nphot, psf_onsky, psf_calib, emccd=self.emccd)
+    

aopera/fiber.py

@@ -0,0 +1,61 @@
+"""
+Step-index single-mode fibers (SMF)
+
+date: 2025-02-15
+author: astriffling, rfetick
+"""
+
+from functools import lru_cache
+import numpy as np
+
+
+def normalized_frequency(r_core:float, na:float, wl:float):
+    """
+    Compute the normalized frequency of the fiber
+
+    Parameters
+    ----------
+    r_core : float, radius of the fiber core [m].
+    na     : float, numerical aperture of the fiber.
+    wl     : float, wavelength of interest [m].
+    """
+    return (2 * np.pi * r_core * na) / wl
+
+
+def mode_field_radius(*args):
+    """
+    Compute the fondamental MFR in [m]
+    """
+    v = normalized_frequency(*args)
+    return args[0] * (0.65 + 1.619/v**(3/2) + 2.879/v**6 - (0.016 + 1.561*v**(-7))) 
+    
+
+@lru_cache(maxsize=1)
+def fundamental_mode(nx:int, sampling:int, *args):
+    """
+    Compute the fundamental mode of the fiber, using gaussian approximation of
+    the mode LP01/TEM00.
+    
+    Parameters
+    ----------
+    nx : int, number of pixel of table.
+    """
+    mfr = mode_field_radius(*args)
+    xx,yy = (np.mgrid[0:nx,0:nx]-nx//2) * args[2] / (2 * args[1] * sampling)
+    return np.exp(-(xx**2 + yy**2) / (mfr**2))
+
+
+def coupling_efficiency(field:np.array, *args):
+    """
+    Compute the SMF coupling efficiency for a focal plane EM field 
+    
+    From G. P. P. L. Otten et al, A&A, 2021, p4
+    https://doi.org/10.1051/0004-6361/202038517
+
+    Parameters
+    ----------
+    field : array.
+    """
+    nx = field.shape[0]
+    mode = fundamental_mode(nx, *args)
+    return np.abs(np.sum(mode * np.conj(field)))**2 / (np.sum(np.abs(mode)**2) * np.sum(np.abs(field)**2))

aopera/otfpsf.py

@@ -0,0 +1,212 @@
+"""
+Convert phase PSD to OTF and to PSF
+"""
+
+
+import numpy as np
+from numpy.fft import fft2, ifft2, ifftshift, fftshift
+from aopera.utils import aperture
+from aopera.readconfig import read_config_file, read_config_tiptop, set_attribute
+from aopera.aopsd import piston_filter, psd_wvl_scaling
+from functools import lru_cache
+
+
+def angle2freq(pix_angle,wvl):
+    """Computes the PSD frequency step [1/m] to get the corresponding angular
+    pixel size [rad] on the PSF."""
+    df = pix_angle/wvl
+    return df
+
+
+@lru_cache(maxsize=2)
+def otf_diffraction(npix, occ=0, samp=2):
+    """
+    Compute the diffraction OTF
+    
+    Parameters
+    ----------
+    npix : int
+        Shape of the output array is (npix,npix).
+    
+    Keywords
+    --------
+    occ : float
+        Eventual central obstruction (0<=occ<1).
+    samp : float
+        Required sampling. Must verify the condition samp>=2.
+    """
+    if samp<2:
+        raise ValueError("Sampling cannot be less than 2.0")
+    aper = aperture(npix, samp=samp, occ=occ)
+    otf = ifftshift(ifft2(np.abs(fft2(aper))**2)) / np.sum(aper)
+    return otf
+
+
+def otf2psf(otf):
+    """Get the PSF for a given OTF"""
+    psf = np.abs(np.real(fftshift(ifft2(fftshift(otf)))))
+    return psf
+
+
+@lru_cache(maxsize=2)
+def psf_diffraction(npix, occ=0, samp=2):
+    """Compute the diffraction PSF"""
+    return otf2psf(otf_diffraction(npix, occ=occ, samp=samp))
+
+
+def psd2otf(psd, df, psdsum=None):
+    """
+    Get the OTF for a given PSD
+    
+    Parameters
+    ----------
+    psd : np.array
+        The PSD 2D array.
+    df : float
+        Frequency step of the PSD array [1/m].
+        
+    Keywords
+    --------
+    psdsum : float, None
+        Integral of the PSD.
+        If psdsum=None, the integral is the numerical sum on the array.
+        The PSF is then of unit energy, since max(otf)=1.
+    """
+    Bphi = ifftshift(np.real(fft2(fftshift(psd)))) * df**2
+    if psdsum is None:
+        psdsum = np.sum(psd) * df**2
+    otf = np.exp(-psdsum+Bphi)
+    return otf
+
+
+def otf_jitter(freq, rms):
+    """
+    Give the Gaussian jitter OTF from its jitter RMS value.
+    Not tested yet, do not use.
+    """
+    return np.exp(-0.5*(2*np.pi*freq*rms)**2)
+
+#%% MODES and DEFORMABLE MIRROR
+def number_modes(nact, occ=0):
+    """Compute the number of corrected modes"""
+    surf_ratio = np.pi*(1-occ**2)/4 #between occulted disk and square
+    nact_total = nact**2
+    return int(round(nact_total * surf_ratio))
+    
+    
+def number_radial(nact, occ=0):
+    """Compute the number of corrected radial modes"""
+    nmode = number_modes(nact, occ=occ)
+    return int(round(np.sqrt(2*nmode)-0.5)) # nr(nr+1)/2 = nmode
+
+
+def controllability(fx, fy, D, actuator, nmode_ratio=1):
+    """Return a DM frequency map that is 1 on the corrected frequencies, 0 otherwise"""
+    cutoff = (actuator-1)/(2*D)
+    if cutoff <= 0: # no DM case
+        return np.zeros(fx.shape)
+    rr = np.sqrt(fx**2+fy**2)
+    if nmode_ratio >1:
+        raise ValueError('nmode_ratio cannot be above 1.')
+    thresh = np.pi/4
+    if nmode_ratio > thresh: # smooth transition from square to circle
+        modal_radial_cutoff = (1-nmode_ratio)/(1-thresh) * (rr <= cutoff)
+        actuator_square_cutoff = (nmode_ratio-thresh)/(1-thresh) * (np.abs(fx)<=cutoff)*(np.abs(fy)<=cutoff)
+        return actuator_square_cutoff + modal_radial_cutoff
+    else:
+        nact_equiv = np.sqrt(nmode_ratio/thresh)
+        return (rr <= (nact_equiv*cutoff)).astype(float)
+
+
+#%% TELESCOPE CLASS
+class Pupil:
+    def __init__(self, dictionary):
+        # self.jitter = 0 # [nm RMS]
+        # self.jitter_angle = 0 # [deg]
+        # self.jitter_freq = 0 # [Hz]
+        set_attribute(self, dictionary, 'pupil')
+    
+    def __repr__(self):
+        s  = 'aopera.PUPIL\n'
+        s += '-------------\n'
+        s += 'diameter   : %.1f m\n'%self.diameter
+        s += 'occultation: %u %%\n'%(100*self.occultation)
+        s += 'nb act     : %u\n'%self.nact
+        s += 'mode ratio : %.2f\n'%self.nmode_ratio
+        s += 'ncpa       : %u nm\n'%self.ncpa
+        # if self.jitter:
+        #     s += 'jitter     : %u nm\n'%self.jitter
+        return s
+    
+    @staticmethod
+    def from_file(filepath, category='pupil'):
+        return Pupil(read_config_file(filepath, category))
+    
+    @staticmethod
+    def from_file_tiptop(filepath):
+        return Pupil(read_config_tiptop(filepath)[0])
+    
+    @staticmethod
+    def from_oopao(tel, dm, M2C):
+        return Pupil({'diameter':tel.D,
+                       'occultation':tel.centralObstruction,
+                       'nact':dm.nAct,
+                       'nmode_ratio':M2C.shape[1]/M2C.shape[0]})
+    
+    @property
+    def surface(self):
+        return np.pi * (self.diameter/2)**2 * (1-self.occultation**2)
+    
+    @property
+    def mode_max(self):
+        return number_modes(self.nact, occ=self.occultation)
+    
+    @property
+    def radial_max(self):
+        return number_radial(self.nact, occ=self.occultation)
+    
+    @property
+    def pitch_dm(self):
+        return self.diameter/(self.nact-1)
+    
+    def pitch_wfs(self, nb_lenslet):
+        return self.diameter/nb_lenslet
+    
+    def pitch_ao(self, nb_lenslet):
+        return max(self.pitch_dm, self.pitch_wfs(nb_lenslet))
+    
+    def cutoff_frequency(self, nb_lenslet):
+        return 1/(2*self.pitch_ao(nb_lenslet))
+    
+    def frequency_step(self, samp):
+        return 1/(samp*self.diameter)
+    
+    def frequency_grid(self, nx, samp):
+        return (np.mgrid[0:nx,0:nx]-nx//2) * self.frequency_step(samp)
+    
+    def piston_filter(self, nx, samp):
+        fx,fy = self.frequency_grid(nx, samp)
+        return piston_filter(np.sqrt(fx**2+fy**2), self.diameter)
+    
+    def otf_diffraction(self, nx, samp):
+        return otf_diffraction(nx, occ=self.occultation, samp=samp)
+    
+    def psf_diffraction(self, nx, samp):
+        return psf_diffraction(nx, occ=self.occultation, samp=samp)
+    
+    def controllability(self, fx, fy):
+        return controllability(fx, fy, self.diameter, self.nact, nmode_ratio=self.nmode_ratio)
+    
+    def otf_ao(self, psd, samp, wvl_scale=1):
+        if wvl_scale == 1:
+            otf_diff = self.otf_diffraction(psd.shape[0], samp)
+            df = self.frequency_step(samp)
+            return psd2otf(psd, df)*otf_diff
+        else:
+            psd_scale = psd_wvl_scaling(psd, 1/wvl_scale)
+            samp_scale = samp * wvl_scale
+            return self.otf_ao(psd_scale, samp_scale)
+    
+    def psf_ao(self, *args, **kwargs):
+        return otf2psf(self.otf_ao(*args, **kwargs))
+