def synlike(a, syntrain, iterations=1000000):
with torch.no_grad():
# note aimag defined with plus here, exp below modified consistently
# [areal] = [aimag] = nbasis x nsignals
areal, aimag = torch.t(a[:, 0::2]), torch.t(a[:, 1::2])
# [anorm] = nsignals
anorm = torch.sum(areal * areal + aimag * aimag, dim=0)
cnt, norm, like = 0, 0, 0
adapt = None
while cnt < iterations:
# [pxt] = nbatch x qdim
_, pxt, alpha = syntrain()
cnt = cnt + alpha.shape[0]
# cpxt = torch.from_numpy(pxt).cuda()
# handle 2D indicator array (assume square qdim)
cpxt = numpy2cuda(
pxt if pxt.ndim == 2 else pxt.reshape(
(pxt.shape[0], pxt.shape[1] * pxt.shape[1])),
alpha.dtype == torch.float32)
# [alphareal] = [alphaimag] = nbatch x nbasis
alphareal, alphaimag = alpha[:, 0::2], alpha[:, 1::2]
# [norm] = qdim
norm += torch.sum(cpxt, dim=0)
# [alphanorm] = nbatch
alphanorm = torch.sum(alphareal * alphareal +
alphaimag * alphaimag,
dim=1)
# automatic normalization of exponentials based on the first batch
loglike = alphareal @ areal + alphaimag @ aimag - 0.5 * alphanorm.unsqueeze(
1) - 0.5 * anorm
if adapt is None:
adapt = torch.max(loglike, dim=0)[0]
loglike -= adapt
# [like] = qdim x nsignals = (qdim x nbatch) @ [(nbatch x nbasis) @ (nbasis x nsignals) + (nbatch x 1) + nsignals]
# remember broadcasting tries to match the last dimension [so A @ b = sum(A * b,axis=1)]
like += torch.t(cpxt) @ torch.exp(loglike)
# (qdim x nsignals) * (qdim x 1)
like = like / norm.unsqueeze(1)
# [ret] = nsignals x qdim = (nsignals x qdim) * qdim
ret = torch.t(like / torch.sum(like, dim=0))
nret = ret.detach().cpu().numpy()
return nret if pxt.ndim == 2 else nret.reshape(
(nret.shape[0], pxt.shape[1], pxt.shape[1]))
def netmeanGn(inputs, net=None, single=True):
if isinstance(inputs, np.ndarray):
inputs = numpy2cuda(inputs, single)
pars = cuda2numpy(net(inputs))
dx = pars[:, 0::3]
std = pars[:, 1::3]
pweight = torch.softmax(torch.from_numpy(pars[:, 2::3]), dim=1).numpy()
# see https://en.wikipedia.org/wiki/Mixture_distribution
xmean = np.sum(pweight * dx, axis=1)
xerr = np.sqrt(np.sum(pweight * (dx**2 + std**2), axis=1) - xmean**2)
return xmean, xerr
def netmeanGn2(inputs, net=None, single=True):
if isinstance(inputs, np.ndarray):
inputs = numpy2cuda(inputs, single)
pars = cuda2numpy(net(inputs))
dx, dy = pars[:, 0::6], pars[:, 2::6]
Fxx, Fyy = pars[:, 1::6]**2, pars[:, 3::6]**2
Fxy = np.arctan(
pars[:, 4::6]) / (0.5 * math.pi) * pars[:, 1::6] * pars[:, 3::6]
det = Fxx * Fyy - Fxy * Fxy
Cxx, Cyy, Cxy = Fyy / det, Fxx / det, -Fxy / det
pweight = torch.softmax(torch.from_numpy(pars[:, 5::6]), dim=1).numpy()
xmean, ymean = np.sum(pweight * dx, axis=1), np.sum(pweight * dy, axis=1)
xerr, yerr = np.sqrt(
np.sum(pweight * (dx**2 + Cxx), axis=1) -
xmean**2), np.sqrt(np.sum(pweight * (dy**2 + Cyy), axis=1) - ymean**2)
xycov = np.sum(pweight * (dx * dy + Cxy), axis=1) - xmean * ymean
return np.vstack((xmean, ymean)).T, np.vstack((xerr, yerr)).T, xycov
def synmean(a, syntrain, iterations=1000000):
with torch.no_grad():
# note aimag defined with plus here, exp below modified consistently
# [areal] = [aimag] = nbasis x nsignals
areal, aimag = torch.t(a[:, 0::2]), torch.t(a[:, 1::2])
# [anorm] = nsignals
anorm = torch.sum(areal * areal + aimag * aimag, dim=0)
cnt = 0
mean, square, cov, norm = 0.0, 0.0, 0.0, 0.0
adapt = None
while cnt < iterations:
# [x] = nbatch
x, _, alpha = syntrain()
cnt = cnt + alpha.shape[0]
x = numpy2cuda(x, alpha.dtype == torch.float32)
# [alphareal] = [alphaimag] = nbatch x nbasis
alphareal, alphaimag = alpha[:, 0::2], alpha[:, 1::2]
# [alphanorm] = nbatch
alphanorm = torch.sum(alphareal * alphareal +
alphaimag * alphaimag,
dim=1)
# [like] = nbatch x nsignals = [(nbatch x nbasis) @ (nbasis x nsignals) + (nbatch x 1) + nsignals]
# like = alphareal @ areal + alphaimag @ aimag - 0.5*alphanorm.unsqueeze(1) - 0.5*anorm
like = alphareal @ areal
like += alphaimag @ aimag
like -= 0.5 * alphanorm.unsqueeze(1)
like -= 0.5 * anorm
if adapt is None:
adapt = torch.max(like, dim=0)[0]
like -= adapt
like = torch.exp_(like) # in-place operation
# [mean] = nsignals = nbatch @ (nbatch x nsignals)
# in the 2D case, 2 x nsignals = 2 x nbatch @ (nbatch x nsignals)
mean += torch.t(x) @ like
square += (torch.t(x)**2) @ like
if x.dim() == 2:
cov += (x[:, 0] * x[:, 1]) @ like
# [norm] = nsignals
norm += torch.sum(like, dim=0)
# note this tensor is very large (iter x len(a)); we should get rid of it asap
del like
# naive variance algorithm, hope for best, see also
# https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance
# [ret] = nsignals
retmean = mean / norm
reterr = torch.sqrt(square / norm - retmean * retmean)
if x.dim() == 1:
return cuda2numpy(retmean).T, cuda2numpy(reterr).T
else:
retcov = cov / norm - retmean[0, :] * retmean[1, :]
return cuda2numpy(retmean).T, cuda2numpy(reterr).T, cuda2numpy(
retcov)
def syntrain(snr=[8, 12],
size=100000,
varx='Mc',
nets=(ar, ai),
seed=None,
noise=1,
varall=False,
region=[[0.2, 0.5], [0.2, 0.25], [-1, 1], [-1, 1]],
single=True):
"""Makes a training set using the ROMAN NN. It returns labels (for `varx`,
or for all if `varall=True`), indicator vectors, and ROM coefficients
(with `snr` and `noise`). Note that the coefficients are kept on the GPU.
Parameters are sampled randomly within `region`."""
device = 'cuda:0' if torch.cuda.is_available() else 'cpu:0'
if seed is not None:
np.random.seed(seed)
torch.manual_seed(seed)
with torch.no_grad():
xs = torch.zeros((size, 4), dtype=torch.float, device=device)
for i, r in enumerate(region):
xs[:, i] = r[0] + (r[1] - r[0]) * torch.rand(
(size, ), dtype=torch.float, device=device)
# handle banks with reduced dimensionality
for i in range(len(region), 4):
xs[:, i] = 0.0
snrs = numpy2cuda(np.random.uniform(*snr, size=size))
alphas = torch.zeros((size, 241 * 2),
dtype=torch.float if single else torch.double,
device=device)
alphar, alphai = nets[0](xs), nets[1](xs)
norm = torch.sqrt(torch.sum(alphar * alphar + alphai * alphai, dim=1))
alphas[:, 0::
2] = snrs[:, np.
newaxis] * alphar / norm[:, np.
newaxis] + noise * torch.randn(
(size, 241),
device=device)
alphas[:, 1::
2] = snrs[:, np.
newaxis] * alphai / norm[:, np.
newaxis] + noise * torch.randn(
(size, 241),
device=device)
xr = np.zeros((size, 5), 'd')
xr[:, :4] = xs.detach().cpu().double().numpy()
xr[:, 4] = snrs.detach().cpu()
del xs, alphar, alphai, norm
# normalize (for provided regions)
for i, r in enumerate(region):
xr[:, i] = (xr[:, i] - r[0]) / (r[1] - r[0])
if isinstance(varx, list):
ix = ['Mc', 'nu', 'chi1', 'chi2'].index(varx[0])
jx = ['Mc', 'nu', 'chi1', 'chi2'].index(varx[1])
i = np.digitize(xr[:, ix], xstops, False) - 1
i[i == -1] = 0
i[i == qdim] = qdim - 1
px = np.zeros((size, qdim), 'd')
px[range(size), i] = 1
j = np.digitize(xr[:, jx], xstops, False) - 1
j[j == -1] = 0
j[j == qdim] = qdim - 1
py = np.zeros((size, qdim), 'd')
py[range(size), j] = 1
if varall:
return xr, np.einsum('ij,ik->ijk', px, py), alphas
else:
return xr[:, [ix, jx]], np.einsum('ij,ik->ijk', px, py), alphas
else:
ix = ['Mc', 'nu', 'chi1', 'chi2'].index(varx)
i = np.digitize(xr[:, ix], xstops, False) - 1
i[i == -1] = 0
i[i == qdim] = qdim - 1
px = np.zeros((size, qdim), 'd')
px[range(size), i] = 1
if varall:
return xr, px, alphas
else:
return xr[:, ix], px, alphas
def syntrainer(net,
syntrain,
lossfunction=None,
iterations=300,
batchsize=None,
initstep=1e-3,
finalv=1e-5,
clipgradient=None,
validation=None,
seed=None,
single=True):
"""Trains network NN against training sets obtained from `syntrain`,
iterating at most `iterations`; stops if the derivative of loss
(averaged over 20 epochs) becomes less than `finalv`."""
if seed is not None:
np.random.seed(seed)
torch.manual_seed(seed)
indicatorloss = 'l' in lossfunction.__annotations__ and lossfunction.__annotations__[
'l'] == 'indicator'
if validation is not None:
raise NotImplementedError
vlabels = numpy2cuda(validation[1] if indicatorloss else validation[0],
single)
vinputs = numpy2cuda(validation[2], single)
optimizer = optim.Adam(net.parameters(), lr=initstep)
training_loss, validation_loss = [], []
for epoch in range(iterations):
t0 = time.time()
xtrue, indicator, inputs = syntrain()
labels = numpy2cuda(indicator if indicatorloss else xtrue, single)
if batchsize is None:
batchsize = inputs.shape[0]
batches = inputs.shape[0] // batchsize
averaged_loss = 0.0
for i in range(batches):
# zero the parameter gradients
optimizer.zero_grad()
# forward + backward + optimize
outputs = net(inputs[i * batchsize:(i + 1) * batchsize])
loss = lossfunction(outputs,
labels[i * batchsize:(i + 1) * batchsize])
loss.backward()
if clipgradient is not None:
torch.nn.utils.clip_grad_norm_(net.parameters(), clipgradient)
optimizer.step()
# print statistics
averaged_loss += loss.item()
training_loss.append(averaged_loss / batches)
if validation is not None:
loss = lossfunction(net(vinputs), vlabels)
validation_loss.append(loss.detach().cpu().item())
if epoch == 1:
print("One epoch = {:.1f} seconds.".format(time.time() - t0))
if epoch % 50 == 0:
print(epoch, training_loss[-1],
validation_loss[-1] if validation is not None else '')
try:
if len(training_loss) > iterations / 10:
training_rate = np.polyfit(range(20),
training_loss[-20:],
deg=1)[0]
if training_rate < 0 and training_rate > -finalv:
print(
f"Terminating at epoch {epoch} because training loss stopped improving sufficiently: rate = {training_rate}"
)
break
if len(validation_loss) > iterations / 10:
validation_rate = np.polyfit(range(20),
validation_loss[-20:],
deg=1)[0]
if validation_rate > 0:
print(
f"Terminating at epoch {epoch} because validation loss started worsening: rate = {validation_rate}"
)
break
except:
pass
print("Final", training_loss[-1],
validation_loss[-1] if validation is not None else '')
if hasattr(net, 'steps'):
net.steps += iterations
else:
net.steps = iterations