point_separation.py 2.66 KB
Newer Older
Jakob Knollmueller's avatar
Jakob Knollmueller committed
1 2
from nifty2go import *
import numpy as np
Jakob Knollmueller's avatar
Jakob Knollmueller committed
3
# from matplotlib import pyplot as plt
Jakob Knollmueller's avatar
Jakob Knollmueller committed
4
from astropy.io import fits
Jakob Knollmueller's avatar
Jakob Knollmueller committed
5 6
from separation_energy import SeparationEnergy
from nifty2go.library.nonlinearities import PositiveTanh
Jakob Knollmueller's avatar
Jakob Knollmueller committed
7

Jakob Knollmueller's avatar
Jakob Knollmueller committed
8
def load_data(path):
Jakob Knollmueller's avatar
Jakob Knollmueller committed
9

Jakob Knollmueller's avatar
Jakob Knollmueller committed
10 11 12 13
    # if path[-5:] == '.fits':
    data = fits.open(path)[1].data
    # else:
    #     data = plt.imread(path)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
14

Jakob Knollmueller's avatar
Jakob Knollmueller committed
15 16 17
    data = data.clip(min=0.001)
    data = np.ndarray.astype(data, float)
    return data
Jakob Knollmueller's avatar
Jakob Knollmueller committed
18

Jakob Knollmueller's avatar
Jakob Knollmueller committed
19
def build_problem(data, alpha):
20
    s_space = RGSpace(data.shape, distances=len(data.shape) * [1])
Jakob Knollmueller's avatar
Jakob Knollmueller committed
21
    data = Field(s_space,val=data)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
22 23 24 25
    FFT = FFTOperator(s_space)
    h_space = FFT.target[0]
    binbounds = PowerSpace.useful_binbounds(h_space, logarithmic = False)
    p_space = PowerSpace(h_space, binbounds=binbounds)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
26 27 28 29 30 31 32
    initial_spectrum = power_analyze(FFT(log(data)), binbounds=p_space.binbounds)
    initial_correlation = create_power_operator(h_space, initial_spectrum)
    initial_x = Field(s_space, val=-1.)
    alpha = Field(s_space, val=alpha)
    q = Field(s_space, val=10e-40)
    pos_tanh = PositiveTanh()
    ICI = GradientNormController(verbose=False, name="ICI",
Jakob Knollmueller's avatar
Jakob Knollmueller committed
33 34 35 36
                                     iteration_limit=500,
                                     tol_abs_gradnorm=1e-5)
    inverter = ConjugateGradient(controller=ICI)

Jakob Knollmueller's avatar
Jakob Knollmueller committed
37 38 39 40 41 42
    parameters = dict(data=data, correlation=initial_correlation,
                      alpha=alpha, q=q,
                      inverter=inverter, FFT=FFT, pos_tanh=pos_tanh)
    separationEnergy = SeparationEnergy(position=initial_x, parameters=parameters)
    return separationEnergy

Jakob Knollmueller's avatar
Jakob Knollmueller committed
43 44
def problem_iteration(energy, iterations=3):
    controller = GradientNormController(verbose=True, tol_abs_gradnorm=0.00000001, iteration_limit=iterations)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
45 46 47 48 49 50 51 52 53 54 55 56
    minimizer = RelaxedNewton(controller=controller)
    energy, convergence = minimizer(energy)
    new_position = energy.position
    h_space = energy.correlation.domain[0]
    FFT = energy.FFT
    binbounds = PowerSpace.useful_binbounds(h_space, logarithmic=False)
    new_power = power_analyze(FFT(energy.s), binbounds=binbounds)
    new_correlation = create_power_operator(h_space, new_power)
    new_parameters = energy.parameters
    new_parameters['correlation'] = new_correlation
    new_energy = SeparationEnergy(new_position, new_parameters)
    return new_energy
Jakob Knollmueller's avatar
Jakob Knollmueller committed
57

58 59 60 61 62 63 64 65 66 67 68 69 70
def build_multi_problem(data, alpha):
    energy_list = []
    for i in range(data.shape[-1]):
        energy = build_problem(data[...,i],alpha)
        energy_list.append(energy)
    return energy_list

def multi_problem_iteration(energy_list):
    new_energy = []
    for energy in energy_list:
        new_energy.append(problem_iteration(energy))
    return new_energy

Jakob Knollmueller's avatar
Jakob Knollmueller committed
71
if __name__ == '__main__':
Jakob Knollmueller's avatar
Jakob Knollmueller committed
72
    pass
Jakob Knollmueller's avatar
Jakob Knollmueller committed
73 74 75