def _ellipsoidLineIntersects_ne(a, b, lineOrigin, lineDirection, directed=True):
lineOrigin = np.require(lineOrigin, dtype=np.float64)
lineDirection = np.require(lineDirection, dtype=np.float64)
# turn into column vectors
direction = lineDirection.T
origin = -lineOrigin[:,None]
radius = np.array([[1/a], [1/a], [1/b]])
directionTimesRadius = ne('direction * radius')
originTimesRadius = ne('origin * radius')
dirDotOri = np.einsum("ij,ij->j", directionTimesRadius, originTimesRadius)
dirDotDir = np.einsum("ij,ij->j", directionTimesRadius, directionTimesRadius)
oriDotOri = np.einsum("ij,ij->j", originTimesRadius, originTimesRadius)
if directed:
if _isInsideEllipsoid(lineOrigin, a, b):
intersects = ne('dirDotOri + sqrt(dirDotOri**2 - oriDotOri*dirDotDir + dirDotDir) >= 0')
else:
intersects = ne('dirDotOri - sqrt(dirDotOri**2 - oriDotOri*dirDotDir + dirDotDir) >= 0')
else:
intersects = ne('dirDotOri**2 - oriDotOri*dirDotDir + dirDotDir >= 0')
return intersects
def dft_nfor_ne(freq, tvec, dvec):
"""
Calculate the Discrete Fourier transform (slow scales with N^2)
The DFT is normalised to have the mean value of the data at zero frequency
:param freq: numpy array, frequency grid calculated from the time vector
:param tvec: numpy array or list, input time(independent) vector, normalised
by the mean of the time vector
:param dvec: numpy array or list, input dependent vector, normalised by the
mean of the data vector
:return wfn: numpy array of complex numbers, spectral window function
:return dft: numpy array of complex numbers, "dirty" discrete Fourier
transform
"""
# -------------------------------------------------------------------------
# Code starts here
# -------------------------------------------------------------------------
# wfn = np.zeros(len(freq), dtype=complex)
# dft = np.zeros(int(len(freq)/2), dtype=complex)
# make vectors in to matrices
fmatrix = np.matrix(freq) # shape (N x 1)
tmatrix = np.matrix(tvec) # shape (M x 1)
# need dvec to be tiled len(freq) times (to enable multiplication)
d_arr = np.tile(dvec, len(freq)).reshape(len(freq), len(dvec))
dmatrix = np.matrix(d_arr) # shape (N x M)
# work out the phase
ftmatrix = fmatrix.T
twopi = 2 * np.pi
# phase = -2*np.pi*fmatrix.T*tmatrix # shape (N x M)
phase = ne('-twopi*ftmatrix*tmatrix')
# work out the phase vector
# phvec = np.cos(phase) + 1j*np.sin(phase) # shape (N x M)
phvec = ne('cos(phase) + 1j*sin(phase)')
# for freq/2 rows
wfn = np.sum(phvec, axis=1) / len(tvec) # shape (N x 1)
# only for the first freq/2 indices
Nfreq = int(len(freq) / 2)
darray = np.array(dmatrix[:Nfreq])
phvecarray = np.array(phvec)[:Nfreq]
# dft = np.sum(np.array(dmatrix)[: Nfreq] * np.array(phvec)[: Nfreq],
# axis=1)/len(tvec) # shape (N/2 x 1)
multiply = ne('darray * phvecarray')
dft = np.sum(multiply, axis=1) / len(tvec)
return wfn, dft
def _ecef2GeodeticOptimized_ne(x, y, z, a=wgs84A, b=wgs84B):
x, y, z = np.asarray(x), np.asarray(y), np.asarray(z)
e2 = (a * a - b * b) / (a * a) # first eccentricity squared
d = (a * a - b * b) / b
p2 = ne('x*x + y*y')
p = ne('sqrt(p2)')
tu = p2
ne('b*z*(1 + d/sqrt(p2 + z*z))/(a*p)', out=tu)
tu2 = ne('tu*tu')
cu3 = ne('(1/sqrt(1 + tu2))**3')
lat = p2
ne('arctan((z + d*cu3*tu2*tu)/(p - e2*a*cu3))', out=lat)
lon = p
ne('arctan2(y,x)', out=lon)
return lat, lon
def dft_ne(freq, tvec, dvec, log=False):
"""
Calculate the Discrete Fourier transform (slow scales with N^2)
The DFT is normalised to have the mean value of the data at zero frequency
:param freq: numpy array, frequency grid calculated from the time vector
:param tvec: numpy array or list, input time(independent) vector, normalised
by the mean of the time vector
:param dvec: numpy array or list, input dependent vector, normalised by the
mean of the data vector
:param log: boolean, if True prints progress to standard output
if False silent
:return wfn: numpy array of complex numbers, spectral window function
:return dft: numpy array of complex numbers, "dirty" discrete Fourier
transform
"""
# -------------------------------------------------------------------------
# Code starts here
# -------------------------------------------------------------------------
wfn = np.zeros(len(freq), dtype=complex)
dft = np.zeros(int(len(freq) / 2), dtype=complex)
# work out all ocnstants before loop
Ntvec, two_pi_t = len(tvec), -2 * np.pi * tvec
Nfreq = int(len(freq) / 2)
# loop around freq
for i in __tqdmlog__(range(len(freq)), log):
freqi = freq[i]
# phase = -2*np.pi*freq[i]*tvec
# phvec = np.array(np.cos(phase) + 1j * np.sin(phase))
phvec = ne('cos(freqi*two_pi_t) + 1j*sin(freqi*two_pi_t)')
#wfn[i] = np.sum(phvec) / len(tvec)
wfn[i] = ne('sum(phvec/Ntvec)')
if i < Nfreq:
# dft[i] = np.sum(dvec*phvec)/len(tvec)
dft[i] = ne('sum(dvec*phvec/Ntvec)')
return wfn, dft
def _ecef2GeodeticOptimized_ne(x, y, z, a=wgs84A, b=wgs84B):
x, y, z = np.asarray(x), np.asarray(y), np.asarray(z)
e2 = (a*a - b*b) / (a*a) # first eccentricity squared
d = (a*a - b*b) / b
p2 = ne('x*x + y*y')
p = ne('sqrt(p2)')
tu = p2
ne('b*z*(1 + d/sqrt(p2 + z*z))/(a*p)', out=tu)
tu2 = ne('tu*tu')
cu3 = ne('(1/sqrt(1 + tu2))**3')
lat = p2
ne('arctan((z + d*cu3*tu2*tu)/(p - e2*a*cu3))', out=lat)
lon = p
ne('arctan2(y,x)', out=lon)
return lat, lon
def _cartesian_to_spherical_ne(x, y, z, with_radius=True):
if with_radius:
xy = ne('x*x + y*y')
r = ne('sqrt(xy + z*z)')
lat = ne('arctan2(z, sqrt(xy))')
lon = xy
ne('arctan2(y, x)', out=lon)
return r, lat, lon
else:
lat = ne('arctan2(z, sqrt(x*x + y*y))')
lon = ne('arctan2(y, x)')
return lat, lon
def _cartesian_to_spherical_ne(x, y, z, with_radius=True):
if with_radius:
xy = ne('x*x + y*y')
r = ne('sqrt(xy + z*z)')
lat = ne('arctan2(z, sqrt(xy))')
lon = xy
ne('arctan2(y, x)', out=lon)
return r, lat, lon
else:
lat = ne('arctan2(z, sqrt(x*x + y*y))')
lon = ne('arctan2(y, x)')
return lat, lon
def _ellipsoidLineIntersection_ne(a, b, lineOrigin, lineDirection, directed=True):
lineOrigin = np.require(lineOrigin, dtype=np.float64)
lineDirection = np.require(lineDirection, dtype=np.float64)
# turn into column vectors
direction = lineDirection.T
origin = -lineOrigin[:,None]
radius = np.array([[1/a], [1/a], [1/b]])
directionTimesRadius = ne('direction * radius')
originTimesRadius = ne('origin * radius')
directionDotOrigin = np.einsum("ij,ij->j", directionTimesRadius, originTimesRadius)
directionDotDirection = np.einsum("ij,ij->j", directionTimesRadius, directionTimesRadius)
originDotOrigin = np.einsum("ij,ij->j", originTimesRadius, originTimesRadius)
root = ne('sqrt(directionDotOrigin**2 - originDotOrigin*directionDotDirection + directionDotDirection)')
if directed:
if _isInsideEllipsoid(lineOrigin, a, b):
d2 = directionDotOrigin
ne('directionDotOrigin + root', out=d2)
dMin = d2
else:
d1 = directionDotOrigin
ne('directionDotOrigin - root', out=d1)
dMin = d1
dMin = _filterPointsOutsideDirectedLine(dMin)
else:
d1 = ne('directionDotOrigin - root')
d2 = directionDotOrigin
ne('directionDotOrigin + root', out=d2)
dMin = _closestDistance(d1, d2)
res = directionTimesRadius
ne('direction * (dMin / directionDotDirection) - origin', out=res)
return res.T
def _spherical_to_cartesian_ne(r, lat, lon, astuple=True):
# using (3,..) instead of (...,3) as shape has better memory access performance
# x, y, and z are then each a contiguous block of memory
res = np.empty((3, ) + lat.shape, lat.dtype)
x = res[0]
y = res[1]
z = res[2]
if r is None:
ne('cos(lat)', out=x)
else:
ne('r * cos(lat)', out=x)
y[:] = x
ne('x * cos(lon)', out=x)
ne('y * sin(lon)', out=y)
if r is None:
ne('sin(lat)', out=z)
else:
ne('r * sin(lat)', out=z)
if astuple:
return x, y, z
else:
res = np.rollaxis(res, 0, res.ndim)
return res
def tan_pix2world(header, px, py, origin, ascartesian=False):
"""
Fast reimplementation of astropy.wcs.wcs.wcs_pix2world with support for
only the TAN projection. Speedup is about 2x.
:rtype: tuple (ra,dec) in degrees, or cartesian coordinates in one array (h,w,3)
if ascartesian=True
"""
assert origin in [0,1]
assert px.shape == py.shape
shape = px.shape
px = px.ravel()
py = py.ravel()
assert header['CTYPE1'] == 'RA---TAN'
assert header['CTYPE2'] == 'DEC--TAN'
assert header['LATPOLE'] == 0.0
lonpole = header['LONPOLE']
ra_ref = header['CRVAL1'] # degrees
dec_ref = header['CRVAL2']
px_ref = header['CRPIX1']
py_ref = header['CRPIX2']
cd = np.array([[header['CD1_1'], header['CD1_2']],
[header['CD2_1'], header['CD2_2']]])
# make pixel coordinates relative to reference point and put them in one array
pxy = np.empty((len(px),2), float)
pxy[:,0] = px
pxy[:,0] -= px_ref
pxy[:,1] = py
pxy[:,1] -= py_ref
if origin == 0:
pxy += 1
# projection plane coordinates
xy = matrix_multiply(cd, pxy[...,np.newaxis]).reshape(pxy.shape) # [18% of execution time]
del pxy
# native spherical coordinates
# spherical projection
if ne:
x = xy[:,0]
y = xy[:,1]
r = ne('sqrt(x*x+y*y)')
else:
r = vectorLengths(xy)
if ne:
lon = ne('arctan2(x, -y)') # [6% of execution time]
else:
lon = np.arctan2(xy[:,0], -xy[:,1])
del xy
#lat = np.arctan(180/(np.pi*r))
# optimized:
with np.errstate(divide='ignore'):
np.reciprocal(r, r)
np.multiply(180/np.pi, r, r)
if ne:
ne('arctan(r)', out=r)
else:
np.arctan(r, r)
lat = r
# celestial spherical coordinates
# spherical rotation
euler_z = ra_ref+90
euler_x = 90-dec_ref
#euler_z2 = lonpole-90
euler_z2 = -(lonpole-90) # for some reason, this needs to be negative, contrary to paper
rotmat = euler_matrix(np.deg2rad(euler_z), np.deg2rad(euler_x), np.deg2rad(euler_z2), 'rzxz')[:3,:3]
lmn = spherical_to_cartesian(None, lat, lon, astuple=False) # [12% of execution time]
lmnrot = matrix_multiply(rotmat, lmn[...,np.newaxis]).reshape(lmn.shape) # [17% of execution time]
if ascartesian:
return lmnrot.reshape(shape + (3,))
dec, ra = cartesian_to_spherical(lmnrot[:,0], lmnrot[:,1], lmnrot[:,2], with_radius=False) # [15% of execution time]
np.rad2deg(dec, dec)
np.rad2deg(ra, ra)
# wrap at 360deg so that values are in [0,360]
ra -= 360
np.mod(ra, 360, ra) # [12% of execution time]
ra = ra.reshape(shape)
dec = dec.reshape(shape)
return ra, dec
def tan_pix2world(header, px, py, origin, ascartesian=False):
"""
Fast reimplementation of astropy.wcs.wcs.wcs_pix2world with support for
only the TAN projection. Speedup is about 2x.
:rtype: tuple (ra,dec) in degrees, or cartesian coordinates in one array (h,w,3)
if ascartesian=True
"""
assert origin in [0, 1]
assert px.shape == py.shape
shape = px.shape
px = px.ravel()
py = py.ravel()
assert header['CTYPE1'] == 'RA---TAN'
assert header['CTYPE2'] == 'DEC--TAN'
assert header['LATPOLE'] == 0.0
lonpole = header['LONPOLE']
ra_ref = header['CRVAL1'] # degrees
dec_ref = header['CRVAL2']
px_ref = header['CRPIX1']
py_ref = header['CRPIX2']
cd = np.array([[header['CD1_1'], header['CD1_2']],
[header['CD2_1'], header['CD2_2']]])
# make pixel coordinates relative to reference point and put them in one array
pxy = np.empty((len(px), 2), float)
pxy[:, 0] = px
pxy[:, 0] -= px_ref
pxy[:, 1] = py
pxy[:, 1] -= py_ref
if origin == 0:
pxy += 1
# projection plane coordinates
xy = matrix_multiply(cd, pxy[..., np.newaxis]).reshape(
pxy.shape) # [18% of execution time]
del pxy
# native spherical coordinates
# spherical projection
if ne:
x = xy[:, 0]
y = xy[:, 1]
r = ne('sqrt(x*x+y*y)')
else:
r = vectorLengths(xy)
if ne:
lon = ne('arctan2(x, -y)') # [6% of execution time]
else:
lon = np.arctan2(xy[:, 0], -xy[:, 1])
del xy
#lat = np.arctan(180/(np.pi*r))
# optimized:
with np.errstate(divide='ignore'):
np.reciprocal(r, r)
np.multiply(180 / np.pi, r, r)
if ne:
ne('arctan(r)', out=r)
else:
np.arctan(r, r)
lat = r
# celestial spherical coordinates
# spherical rotation
euler_z = ra_ref + 90
euler_x = 90 - dec_ref
#euler_z2 = lonpole-90
euler_z2 = -(
lonpole - 90
) # for some reason, this needs to be negative, contrary to paper
rotmat = euler_matrix(np.deg2rad(euler_z), np.deg2rad(euler_x),
np.deg2rad(euler_z2), 'rzxz')[:3, :3]
lmn = spherical_to_cartesian(None, lat, lon,
astuple=False) # [12% of execution time]
lmnrot = matrix_multiply(rotmat, lmn[..., np.newaxis]).reshape(
lmn.shape) # [17% of execution time]
if ascartesian:
return lmnrot.reshape(shape + (3, ))
dec, ra = cartesian_to_spherical(
lmnrot[:, 0], lmnrot[:, 1], lmnrot[:, 2],
with_radius=False) # [15% of execution time]
np.rad2deg(dec, dec)
np.rad2deg(ra, ra)
# wrap at 360deg so that values are in [0,360]
ra -= 360
np.mod(ra, 360, ra) # [12% of execution time]
ra = ra.reshape(shape)
dec = dec.reshape(shape)
return ra, dec
def _spherical_to_cartesian_ne(r, lat, lon, astuple=True):
# using (3,..) instead of (...,3) as shape has better memory access performance
# x, y, and z are then each a contiguous block of memory
res = np.empty((3,) + lat.shape, lat.dtype)
x = res[0]
y = res[1]
z = res[2]
if r is None:
ne('cos(lat)', out=x)
else:
ne('r * cos(lat)', out=x)
y[:] = x
ne('x * cos(lon)', out=x)
ne('y * sin(lon)', out=y)
if r is None:
ne('sin(lat)', out=z)
else:
ne('r * sin(lat)', out=z)
if astuple:
return x,y,z
else:
res = np.rollaxis(res, 0, res.ndim)
return res
