def gradient(X, y, max_steps=100):
N, M = X.shape
firstBacktrackMult = 0.1
nextBacktrackMult = 0.5
armijoMult = 0.1
stepGrowth = 1.25
stepSize = 1.0
recalcRate = 10
backtrackMult = firstBacktrackMult
beta = np.zeros(M)
print('## -f |df/f| |dx/x| step')
print('----------------------------------------------')
for k in range(max_steps):
# Compute the gradient
if k % recalcRate == 0:
Xbeta = X.dot(beta)
eXbeta = da.exp(Xbeta)
func = da.log1p(eXbeta).sum() - y.dot(Xbeta)
e1 = eXbeta + 1.0
gradient = X.T.dot(eXbeta / e1 - y)
steplen = (gradient**2).sum()**0.5
Xgradient = X.dot(gradient)
Xbeta, eXbeta, func, gradient, steplen, Xgradient = da.compute(
Xbeta, eXbeta, func, gradient, steplen, Xgradient)
obeta = beta
oXbeta = Xbeta
# Compute the step size
lf = func
for ii in range(100):
beta = obeta - stepSize * gradient
if ii and np.array_equal(beta, obeta):
stepSize = 0
break
Xbeta = oXbeta - stepSize * Xgradient
# This prevents overflow
if np.all(Xbeta < 700):
eXbeta = np.exp(Xbeta)
func = np.sum(np.log1p(eXbeta)) - np.dot(y, Xbeta)
df = lf - func
if df >= armijoMult * stepSize * steplen**2:
break
stepSize *= backtrackMult
if stepSize == 0:
print('No more progress')
break
df /= max(func, lf)
db = stepSize * steplen / (np.linalg.norm(beta) + stepSize * steplen)
print('%2d %.6e %9.2e %.2e %.1e' % (k + 1, func, df, db, stepSize))
if df < 1e-14:
print('Converged')
break
stepSize *= stepGrowth
backtrackMult = nextBacktrackMult
return beta
def gradient_descent(X, y, max_steps=100, tol=1e-14):
'''Michael Grant's implementation of Gradient Descent.'''
n, p = X.shape
firstBacktrackMult = 0.1
nextBacktrackMult = 0.5
armijoMult = 0.1
stepGrowth = 1.25
stepSize = 1.0
recalcRate = 10
backtrackMult = firstBacktrackMult
beta = np.zeros(p)
y_local = y.compute()
for k in range(max_steps):
# how necessary is this recalculation?
if k % recalcRate == 0:
Xbeta = X.dot(beta)
eXbeta = da.exp(Xbeta)
func = da.log1p(eXbeta).sum() - y.dot(Xbeta)
e1 = eXbeta + 1.0
gradient = X.T.dot(eXbeta / e1 - y)
Xgradient = X.dot(gradient)
Xbeta, eXbeta, func, gradient, Xgradient = da.compute(
Xbeta, eXbeta, func, gradient, Xgradient)
# backtracking line search
lf = func
stepSize, beta, Xbeta, func = compute_stepsize(beta, gradient,
Xbeta, Xgradient,
y_local, func,
**{
'backtrackMult': backtrackMult,
'armijoMult': armijoMult,
'stepSize': stepSize})
if stepSize == 0:
print('No more progress')
break
# necessary for gradient computation
eXbeta = exp(Xbeta)
df = lf - func
df /= max(func, lf)
if df < tol:
print('Converged')
break
stepSize *= stepGrowth
backtrackMult = nextBacktrackMult
return beta
def calc_moments(self):
with h5py.File(self.infile, 'r', rdcc_nbytes=1000 * 1000 * 1000) as f:
data = da.from_array(f['data'],
chunks=(-1, 256, -1, -1)) # CNHW layout
data = da.transpose(data, (1, 2, 3, 0))
dtype = data.dtype
if dtype != np.float32:
print(
'WARNING: data will be saved as float32 but input ist float64!'
)
if self.mean is None:
arr = data
with ProgressBar():
self.mean, self.std = da.compute(arr.mean(axis=[0, 1, 2]),
arr.std(axis=[0, 1, 2]),
num_workers=8)
else:
self.mean, self.std = np.asarray(
self.mean, dtype=dtype), np.asarray(self.std, dtype=dtype)
print('mean: {}, std: {}'.format(list(self.mean), list(self.std)))
if self.log1p_norm:
data_z_norm = (data - self.mean) / self.std
data_log1p = da.sign(data_z_norm) * da.log1p(
da.fabs(data_z_norm))
if self.mean_log1p is None:
arr = data_log1p
with ProgressBar():
self.mean_log1p, self.std_log1p = da.compute(
arr.mean(axis=[0, 1, 2]),
arr.std(axis=[0, 1, 2]),
num_workers=8)
else:
self.mean_log1p, self.std_log1p = np.asarray(
self.mean_log1p,
dtype=dtype), np.asarray(self.std_log1p, dtype=dtype)
print('mean_log1p: {}, std_log1p: {}'.format(
list(self.mean_log1p), list(self.std_log1p)))
def test_arithmetic():
x = np.arange(5).astype('f4') + 2
y = np.arange(5).astype('i8') + 2
z = np.arange(5).astype('i4') + 2
a = da.from_array(x, chunks=(2,))
b = da.from_array(y, chunks=(2,))
c = da.from_array(z, chunks=(2,))
assert eq(a + b, x + y)
assert eq(a * b, x * y)
assert eq(a - b, x - y)
assert eq(a / b, x / y)
assert eq(b & b, y & y)
assert eq(b | b, y | y)
assert eq(b ^ b, y ^ y)
assert eq(a // b, x // y)
assert eq(a ** b, x ** y)
assert eq(a % b, x % y)
assert eq(a > b, x > y)
assert eq(a < b, x < y)
assert eq(a >= b, x >= y)
assert eq(a <= b, x <= y)
assert eq(a == b, x == y)
assert eq(a != b, x != y)
assert eq(a + 2, x + 2)
assert eq(a * 2, x * 2)
assert eq(a - 2, x - 2)
assert eq(a / 2, x / 2)
assert eq(b & True, y & True)
assert eq(b | True, y | True)
assert eq(b ^ True, y ^ True)
assert eq(a // 2, x // 2)
assert eq(a ** 2, x ** 2)
assert eq(a % 2, x % 2)
assert eq(a > 2, x > 2)
assert eq(a < 2, x < 2)
assert eq(a >= 2, x >= 2)
assert eq(a <= 2, x <= 2)
assert eq(a == 2, x == 2)
assert eq(a != 2, x != 2)
assert eq(2 + b, 2 + y)
assert eq(2 * b, 2 * y)
assert eq(2 - b, 2 - y)
assert eq(2 / b, 2 / y)
assert eq(True & b, True & y)
assert eq(True | b, True | y)
assert eq(True ^ b, True ^ y)
assert eq(2 // b, 2 // y)
assert eq(2 ** b, 2 ** y)
assert eq(2 % b, 2 % y)
assert eq(2 > b, 2 > y)
assert eq(2 < b, 2 < y)
assert eq(2 >= b, 2 >= y)
assert eq(2 <= b, 2 <= y)
assert eq(2 == b, 2 == y)
assert eq(2 != b, 2 != y)
assert eq(-a, -x)
assert eq(abs(a), abs(x))
assert eq(~(a == b), ~(x == y))
assert eq(~(a == b), ~(x == y))
assert eq(da.logaddexp(a, b), np.logaddexp(x, y))
assert eq(da.logaddexp2(a, b), np.logaddexp2(x, y))
assert eq(da.exp(b), np.exp(y))
assert eq(da.log(a), np.log(x))
assert eq(da.log10(a), np.log10(x))
assert eq(da.log1p(a), np.log1p(x))
assert eq(da.expm1(b), np.expm1(y))
assert eq(da.sqrt(a), np.sqrt(x))
assert eq(da.square(a), np.square(x))
assert eq(da.sin(a), np.sin(x))
assert eq(da.cos(b), np.cos(y))
assert eq(da.tan(a), np.tan(x))
assert eq(da.arcsin(b/10), np.arcsin(y/10))
assert eq(da.arccos(b/10), np.arccos(y/10))
assert eq(da.arctan(b/10), np.arctan(y/10))
assert eq(da.arctan2(b*10, a), np.arctan2(y*10, x))
assert eq(da.hypot(b, a), np.hypot(y, x))
assert eq(da.sinh(a), np.sinh(x))
assert eq(da.cosh(b), np.cosh(y))
assert eq(da.tanh(a), np.tanh(x))
assert eq(da.arcsinh(b*10), np.arcsinh(y*10))
assert eq(da.arccosh(b*10), np.arccosh(y*10))
assert eq(da.arctanh(b/10), np.arctanh(y/10))
assert eq(da.deg2rad(a), np.deg2rad(x))
assert eq(da.rad2deg(a), np.rad2deg(x))
assert eq(da.logical_and(a < 1, b < 4), np.logical_and(x < 1, y < 4))
assert eq(da.logical_or(a < 1, b < 4), np.logical_or(x < 1, y < 4))
assert eq(da.logical_xor(a < 1, b < 4), np.logical_xor(x < 1, y < 4))
assert eq(da.logical_not(a < 1), np.logical_not(x < 1))
assert eq(da.maximum(a, 5 - a), np.maximum(a, 5 - a))
assert eq(da.minimum(a, 5 - a), np.minimum(a, 5 - a))
assert eq(da.fmax(a, 5 - a), np.fmax(a, 5 - a))
assert eq(da.fmin(a, 5 - a), np.fmin(a, 5 - a))
assert eq(da.isreal(a + 1j * b), np.isreal(x + 1j * y))
assert eq(da.iscomplex(a + 1j * b), np.iscomplex(x + 1j * y))
assert eq(da.isfinite(a), np.isfinite(x))
assert eq(da.isinf(a), np.isinf(x))
assert eq(da.isnan(a), np.isnan(x))
assert eq(da.signbit(a - 3), np.signbit(x - 3))
assert eq(da.copysign(a - 3, b), np.copysign(x - 3, y))
assert eq(da.nextafter(a - 3, b), np.nextafter(x - 3, y))
assert eq(da.ldexp(c, c), np.ldexp(z, z))
assert eq(da.fmod(a * 12, b), np.fmod(x * 12, y))
assert eq(da.floor(a * 0.5), np.floor(x * 0.5))
assert eq(da.ceil(a), np.ceil(x))
assert eq(da.trunc(a / 2), np.trunc(x / 2))
assert eq(da.degrees(b), np.degrees(y))
assert eq(da.radians(a), np.radians(x))
assert eq(da.rint(a + 0.3), np.rint(x + 0.3))
assert eq(da.fix(a - 2.5), np.fix(x - 2.5))
assert eq(da.angle(a + 1j), np.angle(x + 1j))
assert eq(da.real(a + 1j), np.real(x + 1j))
assert eq((a + 1j).real, np.real(x + 1j))
assert eq(da.imag(a + 1j), np.imag(x + 1j))
assert eq((a + 1j).imag, np.imag(x + 1j))
assert eq(da.conj(a + 1j * b), np.conj(x + 1j * y))
assert eq((a + 1j * b).conj(), (x + 1j * y).conj())
assert eq(da.clip(b, 1, 4), np.clip(y, 1, 4))
assert eq(da.fabs(b), np.fabs(y))
assert eq(da.sign(b - 2), np.sign(y - 2))
l1, l2 = da.frexp(a)
r1, r2 = np.frexp(x)
assert eq(l1, r1)
assert eq(l2, r2)
l1, l2 = da.modf(a)
r1, r2 = np.modf(x)
assert eq(l1, r1)
assert eq(l2, r2)
assert eq(da.around(a, -1), np.around(x, -1))
def log1p(A):
return da.log1p(A)
def test_arithmetic():
x = np.arange(5).astype('f4') + 2
y = np.arange(5).astype('i8') + 2
z = np.arange(5).astype('i4') + 2
a = da.from_array(x, chunks=(2, ))
b = da.from_array(y, chunks=(2, ))
c = da.from_array(z, chunks=(2, ))
assert eq(a + b, x + y)
assert eq(a * b, x * y)
assert eq(a - b, x - y)
assert eq(a / b, x / y)
assert eq(b & b, y & y)
assert eq(b | b, y | y)
assert eq(b ^ b, y ^ y)
assert eq(a // b, x // y)
assert eq(a**b, x**y)
assert eq(a % b, x % y)
assert eq(a > b, x > y)
assert eq(a < b, x < y)
assert eq(a >= b, x >= y)
assert eq(a <= b, x <= y)
assert eq(a == b, x == y)
assert eq(a != b, x != y)
assert eq(a + 2, x + 2)
assert eq(a * 2, x * 2)
assert eq(a - 2, x - 2)
assert eq(a / 2, x / 2)
assert eq(b & True, y & True)
assert eq(b | True, y | True)
assert eq(b ^ True, y ^ True)
assert eq(a // 2, x // 2)
assert eq(a**2, x**2)
assert eq(a % 2, x % 2)
assert eq(a > 2, x > 2)
assert eq(a < 2, x < 2)
assert eq(a >= 2, x >= 2)
assert eq(a <= 2, x <= 2)
assert eq(a == 2, x == 2)
assert eq(a != 2, x != 2)
assert eq(2 + b, 2 + y)
assert eq(2 * b, 2 * y)
assert eq(2 - b, 2 - y)
assert eq(2 / b, 2 / y)
assert eq(True & b, True & y)
assert eq(True | b, True | y)
assert eq(True ^ b, True ^ y)
assert eq(2 // b, 2 // y)
assert eq(2**b, 2**y)
assert eq(2 % b, 2 % y)
assert eq(2 > b, 2 > y)
assert eq(2 < b, 2 < y)
assert eq(2 >= b, 2 >= y)
assert eq(2 <= b, 2 <= y)
assert eq(2 == b, 2 == y)
assert eq(2 != b, 2 != y)
assert eq(-a, -x)
assert eq(abs(a), abs(x))
assert eq(~(a == b), ~(x == y))
assert eq(~(a == b), ~(x == y))
assert eq(da.logaddexp(a, b), np.logaddexp(x, y))
assert eq(da.logaddexp2(a, b), np.logaddexp2(x, y))
assert eq(da.exp(b), np.exp(y))
assert eq(da.log(a), np.log(x))
assert eq(da.log10(a), np.log10(x))
assert eq(da.log1p(a), np.log1p(x))
assert eq(da.expm1(b), np.expm1(y))
assert eq(da.sqrt(a), np.sqrt(x))
assert eq(da.square(a), np.square(x))
assert eq(da.sin(a), np.sin(x))
assert eq(da.cos(b), np.cos(y))
assert eq(da.tan(a), np.tan(x))
assert eq(da.arcsin(b / 10), np.arcsin(y / 10))
assert eq(da.arccos(b / 10), np.arccos(y / 10))
assert eq(da.arctan(b / 10), np.arctan(y / 10))
assert eq(da.arctan2(b * 10, a), np.arctan2(y * 10, x))
assert eq(da.hypot(b, a), np.hypot(y, x))
assert eq(da.sinh(a), np.sinh(x))
assert eq(da.cosh(b), np.cosh(y))
assert eq(da.tanh(a), np.tanh(x))
assert eq(da.arcsinh(b * 10), np.arcsinh(y * 10))
assert eq(da.arccosh(b * 10), np.arccosh(y * 10))
assert eq(da.arctanh(b / 10), np.arctanh(y / 10))
assert eq(da.deg2rad(a), np.deg2rad(x))
assert eq(da.rad2deg(a), np.rad2deg(x))
assert eq(da.logical_and(a < 1, b < 4), np.logical_and(x < 1, y < 4))
assert eq(da.logical_or(a < 1, b < 4), np.logical_or(x < 1, y < 4))
assert eq(da.logical_xor(a < 1, b < 4), np.logical_xor(x < 1, y < 4))
assert eq(da.logical_not(a < 1), np.logical_not(x < 1))
assert eq(da.maximum(a, 5 - a), np.maximum(a, 5 - a))
assert eq(da.minimum(a, 5 - a), np.minimum(a, 5 - a))
assert eq(da.fmax(a, 5 - a), np.fmax(a, 5 - a))
assert eq(da.fmin(a, 5 - a), np.fmin(a, 5 - a))
assert eq(da.isreal(a + 1j * b), np.isreal(x + 1j * y))
assert eq(da.iscomplex(a + 1j * b), np.iscomplex(x + 1j * y))
assert eq(da.isfinite(a), np.isfinite(x))
assert eq(da.isinf(a), np.isinf(x))
assert eq(da.isnan(a), np.isnan(x))
assert eq(da.signbit(a - 3), np.signbit(x - 3))
assert eq(da.copysign(a - 3, b), np.copysign(x - 3, y))
assert eq(da.nextafter(a - 3, b), np.nextafter(x - 3, y))
assert eq(da.ldexp(c, c), np.ldexp(z, z))
assert eq(da.fmod(a * 12, b), np.fmod(x * 12, y))
assert eq(da.floor(a * 0.5), np.floor(x * 0.5))
assert eq(da.ceil(a), np.ceil(x))
assert eq(da.trunc(a / 2), np.trunc(x / 2))
assert eq(da.degrees(b), np.degrees(y))
assert eq(da.radians(a), np.radians(x))
assert eq(da.rint(a + 0.3), np.rint(x + 0.3))
assert eq(da.fix(a - 2.5), np.fix(x - 2.5))
assert eq(da.angle(a + 1j), np.angle(x + 1j))
assert eq(da.real(a + 1j), np.real(x + 1j))
assert eq((a + 1j).real, np.real(x + 1j))
assert eq(da.imag(a + 1j), np.imag(x + 1j))
assert eq((a + 1j).imag, np.imag(x + 1j))
assert eq(da.conj(a + 1j * b), np.conj(x + 1j * y))
assert eq((a + 1j * b).conj(), (x + 1j * y).conj())
assert eq(da.clip(b, 1, 4), np.clip(y, 1, 4))
assert eq(da.fabs(b), np.fabs(y))
assert eq(da.sign(b - 2), np.sign(y - 2))
l1, l2 = da.frexp(a)
r1, r2 = np.frexp(x)
assert eq(l1, r1)
assert eq(l2, r2)
l1, l2 = da.modf(a)
r1, r2 = np.modf(x)
assert eq(l1, r1)
assert eq(l2, r2)
assert eq(da.around(a, -1), np.around(x, -1))
def main():
########################################################
infile = '/data/ERA5/hdf5_hr/ds_eval_2000_to_2002.hdf5'
outfile = 'sr_eval_2000_2002.{}.tfrecords'
n_files = 8
gzip = True
shuffle = False
########################################################
SR_ratio = 4
log1p_norm = True
z_norm = False
mean, std = [2.0152406e-08, 2.1581373e-07], [2.8560082e-05, 5.0738556e-05]
mean_log1p, std_log1p = [0.008315503, 0.0028762482], [0.5266841, 0.5418187]
#########################################################
HR_reduce_latitude = 0#107 # 15 deg from each pole
patch_size = None#(96, 96)
n_patches = 50
mode = 'train'
########################################################
f = h5py.File(infile, 'r', rdcc_nbytes=1000*1000*1000)
data = da.from_array(f['data'], chunks=(-1, 16 if shuffle else 256, -1, -1)) # CNHW layout
data = da.transpose(data, (1,2,3,0))
dtype = data.dtype
if dtype != np.float32:
print('WARNING: data will be saved as float32 but input ist float64!')
if mean is None:
arr = data
with ProgressBar():
mean, std = da.compute(arr.mean(axis=[0,1,2]), arr.std(axis=[0,1,2]), num_workers=8)
else:
mean, std = np.asarray(mean, dtype=dtype), np.asarray(std, dtype=dtype)
print('mean: {}, std: {}'.format(list(mean), list(std)))
if log1p_norm:
data_z_norm = (data-mean) / std
data_log1p = da.sign(data_z_norm) * da.log1p(da.fabs(data_z_norm))
if mean_log1p is None:
arr = data_log1p
with ProgressBar():
mean_log1p, std_log1p = da.compute(arr.mean(axis=[0,1,2]), arr.std(axis=[0,1,2]), num_workers=8)
else:
mean_log1p, std_log1p = np.asarray(mean_log1p, dtype=dtype), np.asarray(std_log1p, dtype=dtype)
print('mean_log1p: {}, std_log1p: {}'.format(list(mean_log1p), list(std_log1p)))
data = data_log1p
elif z_norm:
data = (data-mean) / std
if shuffle:
block_indices = np.random.permutation(data.numblocks[0])
else:
block_indices = np.arange(data.numblocks[0])
file_blocks = np.array_split(block_indices, n_files)
i = 0
for n, indices in enumerate(file_blocks):
if n_files > 1:
name = outfile.format(n)
else:
name = outfile
with tf.io.TFRecordWriter(name, options='ZLIB' if gzip else None) as writer:
for block_idx in indices:
block = data.blocks[block_idx].compute()
if shuffle:
block = np.random.permutation(block)
if HR_reduce_latitude:
lat_start = HR_reduce_latitude//2
block = block[:, lat_start:(-lat_start),:, :]
generate_TFRecords(writer, block, SR_ratio, mode, patch_size, n_patches)
i += 1
print('{} / {}'.format(i, data.numblocks[0]), flush=True)
def proximal_grad(X, y, reg='l2', lamduh=0.1, max_steps=100, tol=1e-8):
def l2(x, t):
return 1 / (1 + lamduh * t) * x
def l1(x, t):
return (np.absolute(x) > lamduh * t) * (x - np.sign(x) * lamduh * t)
def identity(x, t):
return x
prox_map = {'l1': l1, 'l2': l2, None: identity}
n, p = X.shape
firstBacktrackMult = 0.1
nextBacktrackMult = 0.5
armijoMult = 0.1
stepGrowth = 1.25
stepSize = 1.0
recalcRate = 10
backtrackMult = firstBacktrackMult
beta = np.zeros(p)
print('# -f |df/f| |dx/x| step')
print('----------------------------------------------')
for k in range(max_steps):
# Compute the gradient
if k % recalcRate == 0:
Xbeta = X.dot(beta)
eXbeta = da.exp(Xbeta)
func = da.log1p(eXbeta).sum() - y.dot(Xbeta)
e1 = eXbeta + 1.0
gradient = X.T.dot(eXbeta / e1 - y)
Xbeta, eXbeta, func, gradient = da.compute(
Xbeta, eXbeta, func, gradient)
obeta = beta
oXbeta = Xbeta
# Compute the step size
lf = func
for ii in range(100):
beta = prox_map[reg](obeta - stepSize * gradient, stepSize)
step = obeta - beta
Xbeta = X.dot(beta).compute() # ugh
# This prevents overflow
if np.all(Xbeta < 700):
eXbeta = np.exp(Xbeta)
func = np.sum(np.log1p(eXbeta)) - np.dot(y, Xbeta)
df = lf - func
if df > 0:
break
stepSize *= backtrackMult
if stepSize == 0:
print('No more progress')
break
df /= max(func, lf)
db = 0
print('%2d %.6e %9.2e %.2e %.1e' % (k + 1, func, df, db, stepSize))
if df < tol:
print('Converged')
break
stepSize *= stepGrowth
backtrackMult = nextBacktrackMult
return beta
