def test_zerns():
n_zerns = 100
nx_size = 128
pxl_scale = 8./nx_size
r0 = 0.2
L0 = 20.
n_scrns = 50000
# Create arrary of zernikes
print("Make Zernikes...")
Zs = aotools.zernikeArray(n_zerns+1, nx_size)[1:]
print("Init phase screen")
phase_screen = infinitephasescreen_fried.PhaseScreen(nx_size, pxl_scale, r0, L0, stencil_length_factor=4)
print("Run tests")
z_coeffs = numpy.zeros((n_scrns, n_zerns))
for i in tqdm(range(n_scrns)):
phase_screen.addRow()
z_coeffs[i] = (phase_screen.scrn * Zs).sum((-1, -2))/(nx_size**2)
z_vars = z_coeffs.var(0)
pyplot.figure()
pyplot.semilogy(z_vars)
pyplot.savefig("infps_N-{}_nz-{}.png".format(nx_size, n_zerns))
pyplot.show()
def test_zerns():
n_zerns = 100
nx_size = 128
pxl_scale = 8. / nx_size
r0 = 0.2
L0 = 20.
n_scrns = 50000
# Create arrary of zernikes
print("Make Zernikes...")
Zs = aotools.zernikeArray(n_zerns + 1, nx_size)[1:]
print("Init phase screen")
phase_screen = infinitephasescreen_fried.PhaseScreen(
nx_size, pxl_scale, r0, L0, stencil_length_factor=4)
print("Run tests")
z_coeffs = numpy.zeros((n_scrns, n_zerns))
for i in tqdm(range(n_scrns)):
phase_screen.addRow()
z_coeffs[i] = (phase_screen.scrn * Zs).sum((-1, -2)) / (nx_size**2)
z_vars = z_coeffs.var(0)
pyplot.figure()
pyplot.semilogy(z_vars)
pyplot.savefig("infps_N-{}_nz-{}.png".format(nx_size, n_zerns))
pyplot.show()
def __init__(self, input_features, output_features):
super(Phase2DLayer, self).__init__()
self.input_features = input_features
self.output_features = output_features
self.z_basis = aotools.zernikeArray(input_features + 1,
output_features,
norm='rms')
self.z_basis = torch.as_tensor(self.z_basis, dtype=torch.float32)
def makeIMatShapes(self):
'''
Creates all the DM shapes which are required for creating the
interaction Matrix. In this case, this is a number of Zernike Polynomials
'''
shapes = aotools.zernikeArray(int(self.n_acts + 1),
int(self.nx_dm_elements))[1:]
pad = self.simConfig.simPad
self.iMatShapes = numpy.pad(shapes, ((0, 0), (pad, pad), (pad, pad)),
mode="constant").astype("float32")
def makeIMatShapes(self):
'''
Creates all the DM shapes which are required for creating the
interaction Matrix. In this case, this is a number of Zernike Polynomials
'''
shapes = aotools.zernikeArray(
int(self.n_acts + 1), int(self.nx_dm_elements))[1:]
pad = self.simConfig.simPad
self.iMatShapes = numpy.pad(
shapes, ((0, 0), (pad, pad), (pad, pad)), mode="constant"
).astype("float32")
def calcInitParams(self, phaseSize=None):
# Call the case WFS methoc
super(Zernike, self).calcInitParams(phaseSize=phaseSize)
# Generate an array of Zernikes to use in the loop
self.n_zerns = self.config.nxSubaps
self.n_measurements = self.n_zerns
# Make Zernike array
self.zernike_array = aotools.zernikeArray(self.n_zerns,
self.pupil_size)
self.zernike_array.shape = self.n_zerns, self.pupil_size**2
mask_sum = self.zernike_array[0].sum()
# Scale to be 1rad in nm ## SHOULDNT NEED THIS AS PHASE ALREADY IN RAD
# self.zernike_array *= ((2 * numpy.pi) / (10e9*self.config.wavelength))
# Normalise for hte number of points in the Zernike anlaysis
self.zernike_array /= mask_sum
def initLGS(self):
"""
Initialises the LGS objects for the WFS
Creates and initialises the LGS objects if the WFS GS is a LGS. This
included calculating the phases additions which are required if the
LGS is elongated based on the depth of the elongation and the launch
position. Note that if the GS is at infinity, elongation is not possible
and a warning is logged.
"""
# Choose the correct LGS object, either with physical or geometric
# or geometric propagation.
if self.lgsConfig.uplink:
lgsObj = eval("LGS.LGS_{}".format(self.lgsConfig.propagationMode))
self.lgs = lgsObj(self.config, self.soapy_config)
else:
self.lgs = None
self.lgsLaunchPos = None
self.elong = 0
self.elongLayers = 0
if self.config.lgs:
self.lgsLaunchPos = self.lgsConfig.launchPosition
# LGS Elongation##############################
if (self.config.GSHeight != 0
and self.lgsConfig.elongationDepth != 0):
self.elong = self.lgsConfig.elongationDepth
self.elongLayers = self.lgsConfig.elongationLayers
# Get Heights of elong layers
self.elongHeights = numpy.linspace(
self.config.GSHeight - self.elong / 2.,
self.config.GSHeight + self.elong / 2., self.elongLayers)
# Calculate the zernikes to add
self.elongZs = aotools.zernikeArray([2, 3, 4], self.pupil_size)
# Calculate the radii of the metapupii at for different elong
# Layer heights
# Also calculate the required phase addition for each layer
self.elongRadii = {}
self.elongPos = {}
self.elongPhaseAdditions = numpy.zeros(
(self.elongLayers, self.los.nx_out_pixels,
self.los.nx_out_pixels))
for i in xrange(self.elongLayers):
self.elongRadii[i] = self.los.findMetaPupilSizes(
float(self.elongHeights[i]))
self.elongPhaseAdditions[i] = self.calcElongPhaseAddition(
i)
self.elongPos[i] = self.calcElongPos(i)
# self.los.metaPupilPos = self.elongPos
logger.debug('Elong Meta Pupil Pos: {}'.format(
self.los.metaPupilPos))
# If GS at infinity cant do elongation
elif (self.config.GSHeight == 0
and self.lgsConfig.elongationDepth != 0):
logger.warning(
"Not able to implement LGS Elongation as GS at infinity")
h = utils.fft(H)
psf = utils.removePadding(np.abs(h) ** 2)
updateAnimation(f, axarr, metrics, phase, utils.removePadding(phaseEst), psf, timer)
return metrics
if __name__ == '__main__':
# Files
reference_file = 'references.fits'
psf_file = 'psf_1.fits'
# Data
wavelength = 2200 * (10**(-9)) #[m]
n=20
z_basis = aotools.zernikeArray(n+1, 128, norm='rms') #[rad]
rv_HDU = fits.open(reference_file)
mask = rv_HDU[0].data # [0-1] function defining entrance pupil
psf_reference = rv_HDU[1].data # diffraction limited point spread function
HDU = fits.open(psf_file)
phase = utils.meterToRadian(HDU[1].data, wavelength* (10**(9)))
H = utils.addPadding(mask) * np.exp(1j * utils.addPadding(phase))
h = utils.fft(H)
psf_test = utils.removePadding(np.abs(h)**2)
metrics = HybridInputOutput(psf_test, mask, phase, n_max=300, animation=True)
metrics = np.array(metrics)
c_zernike[j, i] = c_zernike[j, i] / o_zernike[i]
c_zernike = np.array([
c_zernike[k, :] / np.abs(c_zernike[k, :]).sum() * wavelength * (10**9)
for k in range(n_psfs)
])
# Update scientific camera parameters
config = confParse.loadSoapyConfig(SOAPY_CONF)
config.scis[0].pxlScale = pixelScale
config.tel.telDiam = diameter
config.calcParams()
mask = aotools.circle(config.sim.pupilSize / 2.,
config.sim.simSize).astype(np.float64)
zernike_basis = aotools.zernikeArray(n_zernike + 1,
config.sim.pupilSize,
norm='rms')
psfObj = SCI.PSF(config, nSci=0, mask=mask)
psfs_in = np.zeros(
(n_psfs, psfObj.detector.shape[0], psfObj.detector.shape[1]))
psfs_out = np.zeros(
(n_psfs, psfObj.detector.shape[0], psfObj.detector.shape[1]))
defocus = (wavelength / 4) * (10**9) * zernike_basis[3, :, :]
t0 = time.time()
n_fail = 0
for i in range(n_psfs):
def initLGS(self):
"""
Initialises the LGS objects for the WFS
Creates and initialises the LGS objects if the WFS GS is a LGS. This
included calculating the phases additions which are required if the
LGS is elongated based on the depth of the elongation and the launch
position. Note that if the GS is at infinity, elongation is not possible
and a warning is logged.
"""
# Choose the correct LGS object, either with physical or geometric
# or geometric propagation.
if self.lgsConfig.uplink:
lgsObj = eval("LGS.LGS_{}".format(self.lgsConfig.propagationMode))
self.lgs = lgsObj(self.config, self.soapy_config)
else:
self.lgs = None
self.lgsLaunchPos = None
self.elong = 0
self.elongLayers = 0
if self.config.lgs:
self.lgsLaunchPos = self.lgsConfig.launchPosition
# LGS Elongation##############################
if (self.config.GSHeight!=0 and
self.lgsConfig.elongationDepth!=0):
self.elong = self.lgsConfig.elongationDepth
self.elongLayers = self.lgsConfig.elongationLayers
# Get Heights of elong layers
self.elongHeights = numpy.linspace(
self.config.GSHeight-self.elong/2.,
self.config.GSHeight+self.elong/2.,
self.elongLayers
)
# Calculate the zernikes to add
self.elongZs = aotools.zernikeArray([2,3,4], self.pupil_size)
# Calculate the radii of the metapupii at for different elong
# Layer heights
# Also calculate the required phase addition for each layer
self.elongRadii = {}
self.elongPos = {}
self.elongPhaseAdditions = numpy.zeros(
(self.elongLayers, self.los.nx_out_pixels, self.los.nx_out_pixels))
for i in xrange(self.elongLayers):
self.elongRadii[i] = self.los.findMetaPupilSizes(
float(self.elongHeights[i]))
self.elongPhaseAdditions[i] = self.calcElongPhaseAddition(i)
self.elongPos[i] = self.calcElongPos(i)
# self.los.metaPupilPos = self.elongPos
logger.debug(
'Elong Meta Pupil Pos: {}'.format(self.los.metaPupilPos))
# If GS at infinity cant do elongation
elif (self.config.GSHeight==0 and
self.lgsConfig.elongationDepth!=0):
logger.warning("Not able to implement LGS Elongation as GS at infinity")
def train(model,
dataset,
optimizer,
criterion,
split=[0.9, 0.1],
batch_size=32,
n_epochs=1,
model_dir='./',
random_seed=None,
visdom=False):
# Create directory if doesn't exist
if not os.path.exists(model_dir):
os.makedirs(model_dir)
# Logging
log_path = os.path.join(model_dir, 'logs.log')
utils.set_logger(log_path)
# Visdom support
if visdom:
vis = VisdomWebServer()
# Dataset
dataloaders = {}
dataloaders['train'], dataloaders['val'] = splitDataLoader(
dataset, split=split, batch_size=batch_size, random_seed=random_seed)
# ---
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=150, gamma=0.1)
#scheduler = CosineWithRestarts(optimizer, T_max=40, eta_min=1e-7, last_epoch=-1)
#scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=30, eta_min=1e-7, last_epoch=-1)
# Metrics
metrics_path = os.path.join(model_dir, 'metrics.json')
metrics = {
'model': model_dir,
'optimizer': optimizer.__class__.__name__,
'criterion': criterion.__class__.__name__,
'scheduler': scheduler.__class__.__name__,
'dataset_size': int(len(dataset)),
'train_size': int(split[0] * len(dataset)),
'test_size': int(split[1] * len(dataset)),
'n_epoch': n_epochs,
'batch_size': batch_size,
'learning_rate': [],
'train_loss': [],
'val_loss': [],
'zernike_train_loss': [],
'zernike_val_loss': []
}
# Zernike basis
z_basis = torch.as_tensor(aotools.zernikeArray(100 + 1, 128, norm='rms'),
dtype=torch.float32)
# Device
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# Training
since = time.time()
dataset_size = {
'train': int(split[0] * len(dataset)),
'val': int(split[1] * len(dataset))
}
best_loss = 0.0
for epoch in range(n_epochs):
logging.info('-' * 30)
epoch_time = time.time()
# Each epoch has a training and validation phase
for phase in ['train', 'val']:
if phase == 'train':
model.train() # Set model to training mode
else:
model.eval() # Set model to evaluate mode
running_loss = 0.0
zernike_loss = 0.0
for _, sample in enumerate(dataloaders[phase]):
# GPU support
inputs = sample['image'].to(device)
phase_0 = sample['phase'].to(device)
# zero the parameter gradients
optimizer.zero_grad()
# forward: track history if only in train
with torch.set_grad_enabled(phase == 'train'):
# Network return phase and zernike coeffs
phase_estimation = model(inputs)
loss = criterion(torch.squeeze(phase_estimation), phase_0)
# backward
if phase == 'train':
loss.backward()
optimizer.step()
running_loss += 1 * loss.item() * inputs.size(0)
logging.info('[%i/%i] %s loss: %f' %
(epoch + 1, n_epochs, phase,
running_loss / dataset_size[phase]))
# Update metrics
metrics[phase + '_loss'].append(running_loss / dataset_size[phase])
#metrics['zernike_'+phase+'_loss'].append(zernike_loss / dataset_size[phase])
if phase == 'train':
metrics['learning_rate'].append(get_lr(optimizer))
# Adaptive learning rate
if phase == 'val':
scheduler.step()
# Save weigths
if epoch == 0 or running_loss < best_loss:
best_loss = running_loss
model_path = os.path.join(model_dir, 'model.pth')
torch.save(model.state_dict(), model_path)
# Save metrics
with open(metrics_path, 'w') as f:
json.dump(metrics, f, indent=4)
# Visdom update
if visdom:
vis.update(metrics)
logging.info('[%i/%i] Time: %f s' %
(epoch + 1, n_epochs, time.time() - epoch_time))
time_elapsed = time.time() - since
logging.info('[-----] All epochs completed in {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60))
# L&A cents from off-axis WFS
LAcents = numpy.dot(data[:, :-72], la_mt)
#file = h5py.File('Compare.h5', 'r')
#LAcents = file["sim"][:,:]
########################
# NN centroids
# Read in NN centroids here (set to copy L&A centroids for now)
#NNcents = fits.getdata('CARMEN_datatomo_2013-05-21_12h37m50s_val00.fits')
# NNcents = fits.getdata('your_fits_file_here.fits') # size 10000x72
file = h5py.File('Compare.h5', 'r')
NNcents = file["red"][:, :]
########################
# Calculate wavefronts
zerns = aotools.zernikeArray(zmat.shape[1], 100)
phsTruthTotal = numpy.zeros((100, 100), numpy.float)
phsLATotal = numpy.zeros((100, 100), numpy.float)
phsLAresidual = numpy.zeros((100, 100), numpy.float)
phsNNTotal = numpy.zeros((100, 100), numpy.float)
phsNNresidual = numpy.zeros((100, 100), numpy.float)
TruthWFETotal = 0.
LAWFETotal = 0.
LAresidualWFETotal = 0.
NNWFETotal = 0.
NNresidualWFETotal = 0.
# plot phase or not?
# 1 for plotting
# 0 for not plotting
windSpeedFunction[0] = 0.5*numpy.sin((numpy.arange(nsteps)/float(nsteps))*2*numpy.pi)+1
r0Function[0] = 0.2*numpy.sin((numpy.arange(nsteps)/float(nsteps))*2*numpy.pi)+1
# wind direction is added to original value
windDirectionFunction[0] = 15*numpy.sin((numpy.arange(nsteps)/float(nsteps))*2*numpy.pi)
if debugFunction:
plt.figure('functions')
for ilayer in range(nLayers):
plt.plot(windSpeedFunction[ilayer])
plt.plot(windDirectionFunction[ilayer])
plt.plot(r0Function[ilayer])
plt.draw()
raw_input()
# telescope
zerns = aotools.zernikeArray(nZerns,npix)
pupil = zerns[0]
subapertureMap = aotools.pupil.circle(nsubx/2,nsubx)
subapertureMap -= aotools.pupil.circle(2,nsubx)
# magic WFS
magicSlopes=MagicSHslopes(npix,nsubx,pupil,subapertureMap)
gammas = aotools.makegammas(aotools.zernIndex(nZerns)[0])
subGrads = zernikeSubapSlopes(zerns,gammas[0][:nZerns,:nZerns],gammas[1][:nZerns,:nZerns],nsubx)
z=subGrads.reshape(nZerns,2*nsubx*nsubx)
zmat=numpy.linalg.pinv(z).transpose()
# initiate arrays
phsLayers = numpy.zeros((nLayers,npix,npix),numpy.float)
coeffs = numpy.zeros((zerns.shape[0]),numpy.float)
centroidArray = numpy.zeros((nsteps,2*int(subapertureMap.sum())),numpy.float)