Split catalog fixes jax

DynSpecMS/ClassGiveCatalog.py

@@ -145,14 +145,6 @@ def giveCat(self,SubSet=None):
             self.PosArray=np.asarray(l,dtype=dtype)
             self.DoProperMotionCorr=True
 
-        elif FileCoords is not None:
-            tbl = ascii.read(FileCoords)
-            ra = coord.Angle(tbl["ra"], unit=u.hour)
-            dec = coord.Angle(tbl["dec"], unit=u.degree)
-            self.PosArray=np.zeros((len(tbl),),dtype=dtype)
-            self.PosArray["ra"]=ra.degree
-            self.PosArray["dec"]=dec.degree
-            self.PosArray["Name"][:]=tbl["Name"][:]
         else:
 
             #FileCoords="Transient_LOTTS.csv"
@@ -174,6 +166,16 @@ def giveCat(self,SubSet=None):
         Radius=self.Radius
         self.NOrig=self.PosArray.Name.shape[0]
         Dist=AngDist(self.ra0,self.PosArray.ra,self.dec0,self.PosArray.dec)
+
+        # Log excluded targets
+        ind_out = np.where(Dist >= (Radius*np.pi/180))[0]
+        for i_out in ind_out:
+            name_str = self.PosArray.Name[i_out] if hasattr(self.PosArray, 'Name') else "Target"
+            if isinstance(name_str, bytes):
+                name_str = name_str.decode('utf-8', errors='ignore')
+            dist_deg = Dist[i_out] * 180 / np.pi
+            log.print(f"Skipping {name_str}: distance from phase center ({dist_deg:.3f} deg) > radius ({Radius} deg)")
+
         ind=np.where(Dist<(Radius*np.pi/180))[0]
         self.PosArray=self.PosArray[ind]
 

DynSpecMS/ClassSaveResults.py

@@ -3,15 +3,11 @@
 from __future__ import absolute_import
 from builtins import range
 from builtins import object
-from distutils.spawn import find_executable
 from astropy.time import Time
-from astropy import units as uni
 from astropy.io import fits
 from astropy.wcs import WCS
-from astropy import coordinates as coord
-from astropy import constants as const
 import numpy as np
-import glob, os
+import os
 #import pylab
 from DDFacet.Other import logger
 log=logger.getLogger("ClassSaveResults")
@@ -30,17 +26,14 @@ def GiveMAD(X):
     return np.median(np.abs(X-np.median(X)))
 
 class ClassSaveResults(object):
-    def __init__(self, DynSpecMS,DIRNAME=None):
+    def __init__(self, DynSpecMS, DIRNAME=None):
         self.DynSpecMS=DynSpecMS
-        self.DIRNAME=DIRNAME
-        if self.DIRNAME is None or self.DIRNAME=="MSName":
+
+        # Respect the DIRNAME handed down from the ms2dynspec cleanly
+        if DIRNAME is None or DIRNAME=="MSName":
             self.DIRNAME="DynSpecs_%s"%self.DynSpecMS.OutName
         else:
-            self.DIRNAME=os.path.join(self.DIRNAME,"_DynSpecs_%s"%(self.DynSpecMS.OutName))
-
-
-        #image  = self.DynSpecMS.Image
-        #self.ImageData=np.squeeze(fits.getdata(image, ext=0))
+            self.DIRNAME=DIRNAME
 
         self.ImageI=self.DynSpecMS.ImageI
         if self.ImageI and os.path.isfile(self.DynSpecMS.ImageI):
@@ -178,13 +171,18 @@ def WriteFitsThisDir(self,iDir,Weight=False):
         prihdr.set('CRVAL2', self.DynSpecMS.fMin*1e-6, 'Frequency at the reference pixel (MHz)')
         prihdr.set('CDELT2', self.DynSpecMS.ChanWidth*1e-6, 'Delta freq (MHz)')
         prihdr.set('CUNIT2', 'MHz', 'unit')
-        prihdr.set('CTYPE3', 'Stokes parameter', '1=I, 2=Q, 3=U, 4=V')
+        stokes_labels = []
+        if hasattr(self.DynSpecMS, 'stokes_list'):
+            stokes_labels = [f"{i+1}={s}" for i, s in enumerate(self.DynSpecMS.stokes_list)]
+        else:
+            stokes_labels = ['1=I', '2=Q', '3=U', '4=V']
+        prihdr.set('CTYPE3', 'Stokes parameter', ', '.join(stokes_labels))
         prihdr.set('CRPIX3', 1., 'Reference')
         prihdr.set('CRVAL3', 1., 'frequence at the reference pixel')
         prihdr.set('CDELT3', 1., 'Delta stokes')
         prihdr.set('CUNIT3', '', 'unit')
         prihdr.set('DATE-CRE', Time.now().iso.split()[0], 'Date of file generation')
-        prihdr.set('OBSID', self.DynSpecMS.OutName, 'LOFAR Observation ID')
+        prihdr.set('OBSID', self.DynSpecMS.OutName, 'Observation ID')
         prihdr.set('CHAN-WID', self.DynSpecMS.ChanWidth, 'Frequency channel width')
         prihdr.set('FRQ-MIN', self.DynSpecMS.fMin, 'Minimal frequency')
         prihdr.set('FRQ-MAX', self.DynSpecMS.fMax, 'Maximal frequency')
@@ -193,6 +191,8 @@ def WriteFitsThisDir(self,iDir,Weight=False):
         prihdr.set('RA_RAD', ra, 'Pixel right ascension')
         prihdr.set('DEC_RAD', dec, 'Pixel declination')
         prihdr.set('TEL_NAME', self.DynSpecMS.TELESCOPE_NAME, 'Telescope Name')
+        prihdr.set('PROJECT', self.DynSpecMS.PROJECT, 'Project ID')
+        prihdr.set('OBSERVER', self.DynSpecMS.OBSERVER, 'Observer')
 
         name=self.DynSpecMS.PosArray.Name[iDir]
         if not isinstance(name,str):
@@ -222,7 +222,7 @@ def WriteFitsThisDir(self,iDir,Weight=False):
             # Gn=np.sqrt(Gn0)/Gn1
             # Gn[Gn0==0]=0
         else:
-            Gn = self.DynSpecMS.GOut[iDir,:, :, :].real
+            Gn = self.DynSpecMS.GOut[iDir, :, :, :].real
 
         hdu = fits.PrimaryHDU(np.rollaxis(Gn, 2), header=prihdr)
         #print(f"Fits being written: {fitsname}")
@@ -289,11 +289,6 @@ def PlotSpecSingleDir(self, iDir=0, BoxArcSec=300.):
 
         pylab.clf()
 
-        # if find_executable("latex") is not None:
-        #     pylab.rc('text', usetex=True)
-        # font = {'family':'serif', 'serif': ['Times']}
-        # pylab.rc('font', **font)
-
 
         # Figure properties
         bigfont   = 8

DynSpecMS/MakeDBImagesDynSpec.py

@@ -34,8 +34,6 @@
 log=MyLogger.getLogger("ClassInterpol")
 IdSharedMem=str(int(os.getpid()))+"."
 from pyrap.tables import table
-from killMS.Other.ClassTimeIt import ClassTimeIt
-from killMS.Other.least_squares import least_squares
 import copy
 import astropy.io.fits as pyfits
 import glob

DynSpecMS/dynspecms_version.py

@@ -7,11 +7,17 @@ def version():
     path = os.path.dirname(os.path.abspath(__file__))
     os.chdir(path)
     try:
-        result=subprocess.check_output('git describe --tags', shell=True,universal_newlines=True).rstrip()
-    except:
-        result='unknown'
-    os.chdir(prevdir)
+        result = subprocess.check_output(
+            'git describe --tags', 
+            shell=True,
+            universal_newlines=True,
+            stderr=subprocess.DEVNULL
+        ).rstrip()
+    except Exception:
+        result = 'unknown'
+    finally:
+        os.chdir(prevdir)
     return result
 
 if __name__=='__main__':
-    print(version())
+    print(version())

DynSpecMS/kernels/phase_and_sum_vis.py

@@ -0,0 +1,152 @@
+"""
+JAX-enabled kernel to phase visibilities to directions and sum them.
+
+Behavior:
+- Phases visibilities to a target (ra,dec) given a phase centre (ra0,dec0).
+- Optionally applies Jones (gain) corrections (precomputed or via a callable).
+- Sums visibilities across rows -> returns per-channel, per-polarisation sums
+  and weight sums (and weight-squared sums) compatible with existing code.
+
+Usage sketch:
+  ds, ws, w2s = phase_and_sum_direction(vis, flag, weights,
+                                        u, v, w, A0s, A1s,
+                                        chan_freqs, ra, dec, ra0, dec0,
+                                        slicePol=slice(None),
+                                        Jones=None)
+"""
+from __future__ import division
+import jax.numpy as jnp
+from jax import jit
+JAX_AVAILABLE = True
+from typing import Optional
+from functools import partial
+
+import numpy as np
+
+# speed of light (m/s)
+C = 299792458.0
+
+
+def radec2lm(ra : jnp.ndarray, dec: jnp.ndarray, ra0: jnp.ndarray, dec0: jnp.ndarray) -> tuple[jnp.ndarray, jnp.ndarray]:
+    """
+    Convert RA/DEC to direction cosines l,m for a given phase centre.
+
+    Parameters
+    ----------
+    ra, dec, ra0, dec0 : scalar or array-like (radians)
+        Input coordinates.
+
+    Returns
+    -------
+    l, m : same array type as inputs (jnp or np)
+    """
+    # using the same formula as ClassDynSpecMS.radec2lm
+    l = jnp.cos(dec) * jnp.sin(ra - ra0)
+    m = jnp.sin(dec) * jnp.cos(dec0) - jnp.cos(dec) * jnp.sin(dec0) * jnp.cos(ra - ra0)
+    return l, m
+
+
+@jit
+def _compute_phase(chfreq : jnp.ndarray, u: jnp.ndarray, v: jnp.ndarray, w: jnp.ndarray, l: jnp.ndarray, m: jnp.ndarray, n: jnp.ndarray):
+    """Compute phasor exp(-2pi i nu/c * (u*l + v*m + w*(n-1)))
+
+    Parameters
+    ----------
+    chfreq : array_like, shape (nch,), float
+        Channel centre frequencies in Hz.
+    u, v, w : array_like, shape (nrow, 1, 1)
+        Baseline coordinates in metres.
+    l, m, n : scalar or array-like
+        Direction cosines.
+    Returns
+    -------
+    phase : array_like, shape (nrow, nch, 1), complex
+        Phasor values.
+    """
+    chf = chfreq.reshape((1, -1, 1))
+    kterm = -2.0 * jnp.pi * 1j * chf / C
+    uvw_dot = u * l + v * m + w * (n - 1.0)
+    return jnp.exp(kterm * uvw_dot)
+
+@jit
+def phase_and_sum_direction(vis : jnp.ndarray, flag : jnp.ndarray, weights : jnp.ndarray, u : jnp.ndarray, v : jnp.ndarray, w : jnp.ndarray, A0s : jnp.ndarray, A1s : jnp.ndarray,
+                            chan_freqs : jnp.ndarray, ra : float, dec : float, ra0 : float, dec0 : float,
+                            slicePol=(0,1,2,3), Jones=Optional[dict])-> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]:
+    """
+    Phase visibilities to a direction, apply optional Jones corrections, and sum.
+
+    Parameters
+    ----------
+    vis : array_like, shape (nrow, nch, npol), complex
+        Visibilities for selected time rows.
+    flag : array_like, shape (nrow, nch, npol), bool
+        Flag array (True => flagged).
+    weights : array_like, shape (nrow, nch) or (nrow, nch, 1)
+        Weights (scalar per visibility, broadcasted to pols).
+    u, v, w : array_like, shape (nrow,) or (nrow,1,1)
+        Baseline coordinates in metres.
+    A0s, A1s : array_like, shape (nrow,)
+        Antenna indices per visibility row (used for Jones lookups if provided).
+    chan_freqs : array_like, shape (nch,), float
+        Channel centre frequencies in Hz.
+    ra, dec : float
+        Target direction (radians).
+    ra0, dec0 : float
+        Phase centre (radians).
+    slicePol : slice or sequence
+        Polarisation mapping (like self.slicePol).
+    Jones : tuple of arrays with J0 and J1
+        tuple (J0, J1, ch_mask) : J0,J1 are arrays broadcastable to
+        (nrow, n_ch_mask, 1) and ch_mask selects channels they apply to.
+
+    Returns
+    -------
+    ds : ndarray (nch, noutpol) complex (numpy)
+    ws : ndarray (nch, noutpol) float  (numpy)
+    w2s: ndarray (nch, noutpol) float  (numpy)
+    """
+    nrow, nch, npol = vis.shape
+
+    # build per-pol weights array
+    W = jnp.zeros((nrow, nch, npol), dtype=jnp.float32)
+    # broadcast the (nrow,nch) weights to each pol
+    for i_p in range(npol):
+        W = W.at[:, :, i_p].set(weights)
+
+    # zero weights where flagged
+    W = jnp.where(flag, 0.0, W)
+
+    dcorr = jnp.asarray(vis)
+
+    # compute l,m,n
+    l, m = radec2lm(ra, dec, ra0, dec0)
+    n = jnp.sqrt(jnp.clip(1.0 - l * l - m * m))
+
+    # compute phasing
+    phase = _compute_phase(chan_freqs, u, v, w, l, m, n)
+
+    if Jones is not None:
+        J0, J1, ch_mask = Jones
+
+        dcorr = dcorr.at[:, ch_mask, :].set(J0.conj() * dcorr[:, ch_mask, :] * J1)
+        W = W.at[:, ch_mask, :].set(W[:, ch_mask, :] * (jnp.abs(J0) * jnp.abs(J1)) ** 2)
+
+    # Apply phase and sum across rows
+    dcorr = dcorr * phase
+    ds = jnp.sum(dcorr, axis=0)    # (nch, npol)
+    ws = jnp.sum(W, axis=0)
+    w2s = jnp.sum(W ** 2, axis=0)
+
+    # Apply requested polarisation mapping
+    if isinstance(slicePol, slice):
+        ds_out = ds[:, slicePol]
+        ws_out = ws[:, slicePol]
+        w2s_out = w2s[:, slicePol]
+    else:
+        # assume slicePol is sequence of indices
+        idx = list(slicePol)
+        ds_out = ds[:, idx]
+        ws_out = ws[:, idx]
+        w2s_out = w2s[:, idx]
+
+    return ds_out, ws_out, w2s_out