Setting up inversion for radar

[RemingtonRohel]

Hello, I am struggling to set up a basic inversion for radar measurements. My use case is similar to resolve, but with some key differences. For one, I have 20 static antennas, so the uv plane is quite sparsely sampled compared to the astronomical measurements used in resolve. Also, some of the functions like gridding and w-stacking of resolve are not needed, so I would like to use NIFTy directly.

I would like to model the target using a correlated field model, where the target itself has a complex cross-section to encode phase as well as magnitude. The magnitude is sampled from a log-normal distribution, but I don't know a priori what the hyperparameters of the correlated field model should be. Figure 1 from this paper provides a better explanation of the expected target characteristics.

The measurements are again very similar to those used in resolve, visibilities from all pairs of antennas in the system. For simplicity, I am starting out with a linear array of equally-spaced antennas, which has poor uv sampling but is numerically very simple. I am also constraining the signal space to be a 1D field, again for simplicity.

My biggest issue with implementing using NIFTy is a lack of knowledge of the different domains and operators. There are many options to choose from, so I've been trying to find demos or tutorials that seem close to my use case and copying what is done. However, this has led to RuntimeErrors being raised, so I know I'm a bit out of my depth.

I don't expect a response quickly as Christmas holidays are near, but in the new year any help is greatly appreciated!

For completeness, I've included below the current state of my code (from a Jupyter notebook).

import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline
from collections import OrderedDict

### Set up the antennas and baselines
freq = 10.8e6
wavelength = 299792458 / freq
num_antennas = 16

# [x, y, z] offsets for each antenna
antenna_locations = np.zeros((num_antennas, 3), dtype=np.float32)
antenna_locations[:16, 0] = np.arange(num_antennas) * 15.24  # main array, leftmost is at [0, 0]
# antenna_locations[16:, 0] = np.arange(6, 10) * 15.24  # intf array
# antenna_locations[16:, 1] = -100.0  # intf array 100m behind main array

# [num_antennas, num_antennas]
baselines = antenna_locations[:, np.newaxis] - antenna_locations[np.newaxis, :]
print(baselines.shape)  # should be [16, 16, 3]

### Set up the matrix that maps the sky to the visibilities
directions_deg = np.arange(-30, 30.5)
directions_rad = np.deg2rad(directions_deg)

wave_num = 2.0 * np.pi / wavelength
exponents = np.einsum('ab,c->abc', wave_num * baselines[:, :, 0], np.sin(directions_rad))
B = np.exp(-1j * exponents, dtype=np.complex64)
B = B.reshape((B.shape[0] ** 2, -1))
print(B.shape)  # should be [256, 61]

### My attempt at using NIFTy
N_pixels = len(directions_deg)
correlation_space = ift.RGSpace([num_antennas**2])

class SingleDomain(ift.LinearOperator):
    def __init__(self, domain, target):
        self._domain = ift.makeDomain(domain)
        self._target = ift.makeDomain(target)
        self._capability = self.TIMES | self.ADJOINT_TIMES

    def apply(self, x, mode):
        self._check_input(x, mode)
        return ift.makeField(self._tgt(mode), x.val)

# Set up signal domain
position_space = ift.RGSpace([N_pixels])

# Set up signal model
cfmaker = ift.CorrelatedFieldMaker('')
cfmaker.add_fluctuations(position_space, (0.1, 1e-2), (2, .2), (.01, .5), (-4, 2.),
                         'amp1')
cfmaker.set_amplitude_total_offset(0., (1e-2, 1e-6))
correlated_field = cfmaker.finalize()
# correlated_field = ift.SimpleCorrelatedField(

normalized_amp = cfmaker.get_normalized_amplitudes()
pspec = normalized_amp[0]**2
DC = SingleDomain(correlated_field.target, position_space)

# Apply a nonlinearity
signal = DC @ ift.exp(correlated_field)

# Build the signal response
M = ift.DiagonalOperator(ift.Field.full(ift.RGSpace(N_pixels), 1.))

class CorrelationOperator(ift.LinearOperator):
    def __init__(self):
        self._domain = ift.DomainTuple.make(M.domain)
        self._target = ift.DomainTuple.make(correlation_space)
        self._B = B
        self._B_inv = np.linalg.pinv(B)
        self._B_adjoint_inv = np.linalg.pinv(B.T.conj())
        self._capability = self.TIMES | self.ADJOINT_TIMES | self.INVERSE_TIMES
        
    def apply(self, x, mode):
        self._check_input(x, mode)
        if mode == self.TIMES:
            y = self._B @ x.val
        elif mode == self.ADJOINT_TIMES:
            y = self._B.T.conj() @ x.val
        elif mode == self.INVERSE_TIMES:
            y = self._B_inv @ x.val
        elif mode == self.ADJOINT_INVERSE_TIMES:
            y = self._B_adjoint_inv @ x.val
        else:
            raise NotImplementedError
        return ift.makeField(self._tgt(mode), y)
            
correlator = CorrelationOperator()
GR = ift.GeometryRemover(correlator.target)
R = GR(correlator(M))
signal_response = R(signal)

# Specify noise
data_space = R.target
noise = 0.001
print(data_space)
mock_noise = ift.Field.full(data_space, noise)
N = ift.DiagonalOperator(mock_noise, data_space, sampling_dtype=complex)

# Generate mock signal and data
mock_position = ift.from_random(signal_response.domain, 'normal', dtype=complex)
data = signal_response(mock_position) + N.draw_sample()

plot = ift.Plot()
plot.add(signal(mock_position), title='Ground Truth')
plot.add(R.adjoint_times(data), title='Data')
plot.add([pspec.force(mock_position)], title='Power Spectrum')

# Minimization parameters
ic_sampling = ift.AbsDeltaEnergyController(name='Sampling', deltaE=0.01, iteration_limit=100)
ic_newton = ift.AbsDeltaEnergyController(name='Newton', deltaE=0.01, iteration_limit=35)
minimizer = ift.NewtonCG(ic_newton, enable_logging=True)

# Set up likelihood energy and information Hamiltonian
likelihood_energy = ift.GaussianEnergy(data, sampling_dtype=float) @ signal_response

n_samples = 20
n_iterations = 5
samples = ift.optimize_kl(likelihood_energy, n_iterations, n_samples,
                          minimizer, ic_sampling, None, comm=comm,
                          # output_directory="getting_started_5_results",
                          # export_operator_outputs={"signal": signal, "power spectrum": pspec}
                         )

# Plotting
filename_res = filename.format("results")
plot = ift.Plot()
mean, var = samples.sample_stat(signal)
plot.add(mean, title="Posterior Mean")
plot.add(ift.sqrt(var), title="Posterior Standard Deviation")

n_samples = samples.n_samples
plot.add(list(samples.iterator(pspec)) + [samples.average(pspec.log()).exp(),
          pspec.force(mock_position)],
         title="Sampled Posterior Power Spectrum",
         linewidth=[1.]*n_samples + [3., 3.])
plot.show()

Running this yields the following traceback:

An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.
Offset amplitude: 1.00E-02 ± 7.77E-07
Total fluctuation amplitude: 9.99E-02 ± 9.98E-03
/home/remington/repos/pp/.venv/lib64/python3.9/site-packages/nifty8/library/correlated_fields.py:144: ComplexWarning: Casting complex values to real discards the imaginary part
  res[from_third] = np.cumsum(x[second], axis=axis)
/home/remington/repos/pp/.venv/lib64/python3.9/site-packages/nifty8/library/correlated_fields.py:145: ComplexWarning: Casting complex values to real discards the imaginary part
  res[from_third] = (res[from_third] + res[no_border])/2*self._log_vol[extender_sl] + x[first]
DomainTuple, len: 1
* UnstructuredDomain(shape=(256,))
---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
File ~/repos/pp/.venv/lib64/python3.9/site-packages/nifty8/operators/sampling_enabler.py:66, in SamplingEnabler.special_draw_sample(self, from_inverse)
     65 try:
---> 66     res = self._op.draw_sample(from_inverse)
     67     return self._op(res), res

File ~/repos/pp/.venv/lib64/python3.9/site-packages/nifty8/operators/sum_operator.py:225, in SumOperator.draw_sample(self, from_inverse)
    224 if from_inverse:
--> 225     raise NotImplementedError(
    226         "cannot draw from inverse of this operator")
    227 res = None

NotImplementedError: cannot draw from inverse of this operator

During handling of the above exception, another exception occurred:

UFuncTypeError                            Traceback (most recent call last)
Cell In[5], line 100
     98 n_samples = 20
     99 n_iterations = 5
--> 100 samples = ift.optimize_kl(likelihood_energy, n_iterations, n_samples,
    101                           minimizer, ic_sampling, None, comm=comm,
    102                           # output_directory="getting_started_5_results",
    103                           # export_operator_outputs={"signal": signal, "power spectrum": pspec}
    104                          )
    106 # Plotting
    107 filename_res = filename.format("results")

File ~/repos/pp/.venv/lib64/python3.9/site-packages/nifty8/minimization/optimize_kl.py:395, in optimize_kl(likelihood_energy, total_iterations, n_samples, kl_minimizer, sampling_iteration_controller, nonlinear_sampling_minimizer, constants, point_estimates, transitions, export_operator_outputs, output_directory, initial_position, initial_index, comm, inspect_callback, terminate_callback, plot_energy_history, plot_minisanity_history, save_strategy, return_final_position, resume, sanity_checks)
    393         mean = comm(iglobal).bcast(mean, root=0)
    394 else:
--> 395     e = SampledKLEnergy(
    396         mean_iter,
    397         ham,
    398         n_samples(iglobal),
    399         nonlinear_sampling_minimizer(iglobal),
    400         comm=comm(iglobal),
    401         constants=constants(iglobal),
    402         point_estimates=point_estimates(iglobal))
    403     e, _ = minimizer(e)
    404     mean = MultiField.union([mean, e.position]) if mf_dom else e.position

File ~/repos/pp/.venv/lib64/python3.9/site-packages/nifty8/minimization/kl_energies.py:290, in SampledKLEnergy(position, hamiltonian, n_samples, minimizer_sampling, mirror_samples, constants, point_estimates, napprox, comm, nanisinf)
    287 position, hamiltonian = _reduce_by_keys(position, hamiltonian, invariant)
    289 _, ham_sampling = _reduce_by_keys(position, hamiltonian, point_estimates)
--> 290 sample_list = draw_samples(position, ham_sampling, minimizer_sampling, n_samples,
    291                            mirror_samples, napprox=napprox, comm=comm)
    292 return SampledKLEnergyClass(sample_list, hamiltonian, constants, inv_pos, nanisinf)

File ~/repos/pp/.venv/lib64/python3.9/site-packages/nifty8/minimization/kl_energies.py:145, in draw_samples(position, H, minimizer, n_samples, mirror_samples, napprox, want_error, comm)
    143 neg = mirror_samples and (i % 2 != 0)
    144 if not neg or y is None:  # we really need to draw a sample
--> 145     y, yi = met.special_draw_sample(True)
    147 if geometric:
    148     m = transformation_mean - y if neg else transformation_mean + y

File ~/repos/pp/.venv/lib64/python3.9/site-packages/nifty8/operators/sampling_enabler.py:76, in SamplingEnabler.special_draw_sample(self, from_inverse)
     74 else:
     75     s = self._prior.draw_sample(from_inverse=True)
---> 76     nj = self._likelihood.draw_sample()
     77     b = self._prior(s) + nj
     78     energy = QuadraticEnergy(s, self._op, b,
     79                              _grad=self._likelihood(s) - nj)

File ~/repos/pp/.venv/lib64/python3.9/site-packages/nifty8/operators/sandwich_operator.py:94, in SandwichOperator.draw_sample(self, from_inverse)
     90     raise NotImplementedError(
     91         "cannot draw from inverse of this operator")
     93 # Samples from general sandwiches
---> 94 return self._bun.adjoint_times(
     95     self._cheese.draw_sample(from_inverse))

File ~/repos/pp/.venv/lib64/python3.9/site-packages/nifty8/operators/linear_operator.py:221, in LinearOperator.adjoint_times(self, x)
    208 def adjoint_times(self, x):
    209     """Applies the adjoint-Operator to a given Field.
    210 
    211     Parameters
   (...)
    219         The processed Field defined on the Operator's domain.
    220     """
--> 221     return self.apply(x, self.ADJOINT_TIMES)

File ~/repos/pp/.venv/lib64/python3.9/site-packages/nifty8/operators/chain_operator.py:148, in ChainOperator.apply(self, x, mode)
    146 t_ops = self._ops if mode & self._backwards else reversed(self._ops)
    147 for op in t_ops:
--> 148     x = op.apply(x, mode)
    149 return x

File ~/repos/pp/.venv/lib64/python3.9/site-packages/nifty8/operators/sum_operator.py:216, in SumOperator.apply(self, x, mode)
    214 res = None
    215 for op, neg in zip(self._ops, self._neg):
--> 216     tmp = op.apply(x.extract(op._dom(mode)), mode)
    217     if res is None:
    218         res = -tmp if neg else tmp

File ~/repos/pp/.venv/lib64/python3.9/site-packages/nifty8/operators/chain_operator.py:148, in ChainOperator.apply(self, x, mode)
    146 t_ops = self._ops if mode & self._backwards else reversed(self._ops)
    147 for op in t_ops:
--> 148     x = op.apply(x, mode)
    149 return x

File ~/repos/pp/.venv/lib64/python3.9/site-packages/nifty8/operators/sum_operator.py:216, in SumOperator.apply(self, x, mode)
    214 res = None
    215 for op, neg in zip(self._ops, self._neg):
--> 216     tmp = op.apply(x.extract(op._dom(mode)), mode)
    217     if res is None:
    218         res = -tmp if neg else tmp

    [... skipping similar frames: ChainOperator.apply at line 148 (3 times), SumOperator.apply at line 216 (2 times)]

File ~/repos/pp/.venv/lib64/python3.9/site-packages/nifty8/operators/sum_operator.py:216, in SumOperator.apply(self, x, mode)
    214 res = None
    215 for op, neg in zip(self._ops, self._neg):
--> 216     tmp = op.apply(x.extract(op._dom(mode)), mode)
    217     if res is None:
    218         res = -tmp if neg else tmp

File ~/repos/pp/.venv/lib64/python3.9/site-packages/nifty8/operators/chain_operator.py:148, in ChainOperator.apply(self, x, mode)
    146 t_ops = self._ops if mode & self._backwards else reversed(self._ops)
    147 for op in t_ops:
--> 148     x = op.apply(x, mode)
    149 return x

File ~/repos/pp/.venv/lib64/python3.9/site-packages/nifty8/library/correlated_fields.py:151, in _TwoLogIntegrations.apply(self, x, mode)
    149 res = np.zeros(self._domain.shape)
    150 x[from_third] = np.cumsum(x[from_third][reverse], axis=axis)[reverse]
--> 151 res[first] += x[from_third]
    152 x[from_third] *= (self._log_vol/2.)[extender_sl]
    153 x[no_border] += x[from_third]

UFuncTypeError: Cannot cast ufunc 'add' output from dtype('complex128') to dtype('float64') with casting rule 'same_kind'