import torch
from data_management import load_dataset
from networks import IterativeNet, Tiramisu
from operators import TVAnalysis, get_tikhonov_matrix
# --- load configuration -----
import config # isort:skip
# ----- general setup -----
mpl.use("agg")
device = torch.device("cuda:0")
# ----- operators -----
OpA = config.meas_op(config.m, config.n, device=device, **config.meas_params)
OpTV = TVAnalysis(config.n, device=device)
# ----- build linear inverter ------
reg_fac = 2e-2
inverter = torch.nn.Linear(OpA.m, OpA.n, bias=False)
inverter.weight.requires_grad = False
inverter.weight.data = get_tikhonov_matrix(OpA, OpTV, reg_fac)
# ----- network configuration -----
subnet_params = {
"in_channels": 1,
"out_channels": 1,
"drop_factor": 0.0,
"down_blocks": (5, 7, 9, 12, 15),
"up_blocks": (15, 12, 9, 7, 5),
def sample_tv_signal(
n,
j_min=10,
j_max=20,
min_dist=5,
bound=5,
min_height=0.2,
n_seed=None,
t_seed=None,
):
""" Creates a random piecewise constant signal.
Creates a piecewise constant signal of shape (n,) with a random number of
"jumps" (discontinuities).
With minimal jump size and zero boundaries.
Parameters
----------
n : int
Size of the signal.
j_min : int, optional
Minimum number of jumps in the signal. (Default 10)
j_max : int, optional
Maximum number of jumps in teh signal (Default 20)
min_dist : int, optional
Minimum distance between the jump positions. (Default 5)
bound : int, optional
Pads each side with boundary zeros.
min_height : double, optional
Minimal jump height of jumps in the interior.
n_seed : int, optional
Seed for the numpy random number generator for drawing the jump
positions. Set to `None` for not setting any seed and keeping the
current state of the random number generator. (Default `None`)
t_seed : int, optional
Seed for the troch random number generator for drawing the jump
heights. Set to `None` for not setting any seed and keeping the
current state of the random number generator. (Default `None`)
Returns
-------
torch.Tensor
Will be of shape (n,).
"""
if n_seed is not None:
np.random.seed(n_seed)
if t_seed is not None:
torch.manual_seed(t_seed)
# restricted dimension
n_restr = n - 2 * bound
# check if jump distance condition can be satisfied
if (n_restr) / (min_dist) <= (j_max + 1):
raise ValueError(
"The signal size {} is not large enough to contain {} "
"jumps with minimum distance {}.".format(n_restr, j_max, min_dist))
# draw the number of jumps
num = np.random.randint(j_min - 2, j_max - 1) # +2 jumps whp for 0s
# draw the jump positions
pos = np.sort(
np.random.choice(
np.arange(min_dist, n_restr - min_dist - (num - 1) *
(min_dist - 1)),
num,
False,
)) + np.arange(num) * (min_dist - 1)
# synthesize signal
coefs = torch.zeros(n_restr)
coefs[-1] = torch.randn(1) # random signal mean (constant shift)
coefs[pos] = torch.randn(num) # random jump heights around the mean
coefs[torch.abs(coefs) < min_height] = (
torch.sign(coefs[torch.abs(coefs) < min_height]) * min_height
) # minimal height
x_restr = TVSynthesis(n_restr)(coefs)
# pad with 0s
x = torch.zeros(n)
x[bound:n_restr + bound] = x_restr
# real coefs
coefs = TVAnalysis(n)(x)
return x, coefs