forked from idia-astro/image-generator
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathmake_image.py
executable file
·187 lines (146 loc) · 8.71 KB
/
make_image.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
#!/usr/bin/env python3
# Script for generating synthetic FITS files
# Recipe for writing files too large to fit in memory taken from:
# https://docs.astropy.org/en/stable/generated/examples/io/skip_create-large-fits.html
import argparse
import itertools
import sys
import numpy as np
from astropy.io import fits
from astropy.modeling.models import Gaussian2D
from astropy.stats import gaussian_fwhm_to_sigma
NAN_OPTIONS = ("pixel", "row", "column", "channel", "stokes", "image")
INF_OPTIONS = ("positive", "negative")
def make_image(args):
dims = tuple(args.dimensions)
N = len(dims)
# create header
dummy_dims = tuple(1 for d in dims)
dummy_data = np.zeros(dummy_dims, dtype=np.float32)
hdu = fits.PrimaryHDU(data=dummy_data)
header = hdu.header
for i, dim in enumerate(dims, 1):
header["NAXIS%d" % i] = dim
if args.header:
extraheader = fits.Header.fromstring(args.header, sep="\n")
for card in extraheader.cards:
header.append(card)
header.tofile(args.output, overwrite=True)
# create full-sized zero image
header_size = len(header.tostring()) # Probably 2880. We don't pad the header any more; it's just the bare minumum
data_size = (np.prod(dims) * np.dtype(np.float32).itemsize)
# This is not documented in the example, but appears to be Astropy's default behaviour
# Pad the total file size to a multiple of the header block size
block_size = 2880
data_size = block_size * (((data_size - 1)//block_size) + 1)
with open(args.output, "rb+") as f:
f.seek(header_size + data_size - 1)
f.write(b"\0")
# write random data
# optionally seed the random number generator
if args.seed is not None:
rng = np.random.default_rng(args.seed)
else:
rng = np.random.default_rng()
hdul = fits.open(args.output, "update", memmap=True)
data = hdul[0].data
if not args.max_bytes:
strip_size = np.prod(data.shape[-2:])
else:
strip_size = args.max_bytes // np.dtype(np.float32).itemsize
total_size = np.prod(dims)
remainder = total_size % strip_size
rounded_size = total_size - remainder
contiguous_data = data.ravel()
width, height = dims[0:2]
channel_size = width * height
depth = dims[2] if N > 2 else 1
stokes = dims[3] if N > 3 else 1
shaped_data = data.reshape(stokes, depth, height, width)
def nan_sample(start, end):
nan_sample_size = int(np.round((args.nan_density/100) * (end - start)))
return rng.choice(range(start, end), nan_sample_size, False).astype(int)
def inf_sample(start, end):
inf_sample_size = int(np.round((args.inf_density/100) * (end - start)))
return rng.choice(range(start, end), inf_sample_size, False).astype(int)
inf_choice = []
if "positive" in args.infs:
inf_choice.append(np.inf)
if "negative" in args.infs:
inf_choice.append(-np.inf)
if "image" in args.nans:
for i in range(0, rounded_size, strip_size):
contiguous_data[i:i+strip_size] = np.nan
contiguous_data[rounded_size:total_size] = np.nan
elif args.checkerboard:
ch = args.checkerboard
for i in np.arange(0, width, ch * 2):
for j in np.arange(0, height, ch * 2):
shaped_data[:, :, i:i+ch, j:j+ch] = 0.75
shaped_data[:, :, i+ch:i+ch*2, j+ch:j+ch*2] = 0.75
shaped_data[:, :, i:i+ch, j+ch:j+ch*2] = 0.5
shaped_data[:, :, i+ch:i+ch*2, j:j+ch] = 0.5
else:
for i in range(0, rounded_size, strip_size):
contiguous_data[i:i+strip_size] = rng.normal(size=strip_size).astype(np.float32)
if "pixel" in args.nans:
contiguous_data[nan_sample(i, i + strip_size)] = np.nan
if args.infs:
sample = inf_sample(i, i + strip_size)
contiguous_data[sample] = rng.choice(inf_choice, sample.size)
if remainder:
contiguous_data[rounded_size:total_size] = rng.normal(size=remainder).astype(np.float32)
if "pixel" in args.nans:
contiguous_data[nan_sample(rounded_size, total_size)] = np.nan
if args.infs:
sample = inf_sample(rounded_size, total_size)
contiguous_data[sample] = rng.choice(inf_choice, sample.size)
# add row, column, channel and stokes nans here in separate passes
# For now an implementation which ignores max_bytes
if "stokes" in args.nans:
nan_stokes = nan_sample(0, stokes)
for s in nan_stokes:
shaped_data[s] = np.nan
if "channel" in args.nans:
nan_channels = nan_sample(0, depth * stokes)
for channel in nan_channels:
s, c = divmod(channel, depth)
shaped_data[s,c] = np.nan
if "row" in args.nans:
nan_rows = nan_sample(0, height * depth * stokes)
for row in nan_rows:
s, row = divmod(row, height * depth)
c, y = divmod(row, height)
shaped_data[s,c,y] = np.nan
if "column" in args.nans:
nan_columns = nan_sample(0, width * depth * stokes)
for column in nan_columns:
s, column = divmod(column, width * depth)
c, x = divmod(column, width)
shaped_data[s,c,:,x] = np.nan
if args.gaussian_model and (len(args.gaussian_model) - 1) / 6 == int(args.gaussian_model[0]):
params = args.gaussian_model
for k in np.arange(int(params[0])):
for i in np.arange(width):
for j in np.arange(height):
shaped_data[:, :, i, j] += Gaussian2D(params[k * 6 + 3], params[k * 6 + 1], params[k * 6 + 2], params[k * 6 + 4] * gaussian_fwhm_to_sigma, params[k * 6 + 5] * gaussian_fwhm_to_sigma, params[k * 6 + 6] * np.pi / 180)(j, i)
hdul.close()
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Generate a synthetic FITS file.")
parser.add_argument("dimensions", metavar="N", type=int, nargs="+", help="The dimensions of the file, in order (XY, XYZ or XYZW), separated by spaces.")
parser.add_argument("-m", "--max-bytes", type=int, help="The maximum size of image data (in bytes) to create in memory at once. Default is the channel size.")
parser.add_argument("-n", "--nans", nargs="+", help="Options for inserting random NaNs, which are cumulative, separated by spaces. Any combination of %s. By default no NaNs are inserted. 'image' overrides all other options and creates an image full of NaNs. Channel, row and column algorithm currently ignores the --max-bytes restriction." % ", ".join(repr(o) for o in NAN_OPTIONS), default=[])
parser.add_argument("-d", "--nan-density", type=float, help="The density of NaNs to insert, as a percentage. Default: 25. Ignored if -n/--nans is unset.", default=25.0)
parser.add_argument("-i", "--infs", nargs="+", help="Options for inserting random INF pixels, which are cumulative, separated by spaces. Any combination of %s. By default no INFs are inserted." % ", ".join(repr(o) for o in INF_OPTIONS), default=[])
parser.add_argument("--inf-density", type=float, help="The density of INFs to insert, as a percentage. Default: 1. Ignored if -i/--infs is unset.", default=1.0)
parser.add_argument("-c", "--checkerboard", type=int, help="If this is set, the image is filled with a checkerboard pattern with each square the specified number of pixels. If it is unset, Gaussian noise is used. The checkerboard algorithm currently ignores the --max-bytes restriction. This option can be used together with the NaN options.")
parser.add_argument("--gaussian-model", type=float, nargs="+", help="If this is set with correct number of parameters (6n + 1), the image is filled with a Gaussian model in addition to Gaussian noise. A list of float including number of Gaussian components (n), center x (px), center y (px), amplitude, fwhm x (px), fwhm y (px), and p.a. (deg) of component 1, center x of component 2, etc.")
parser.add_argument("-o", "--output", help="The output file name.")
parser.add_argument("-s", "--seed", help="Seed for the random number generator.", type=int)
parser.add_argument("-H", "--header", help="Additional entries to append to the header, as a literal string with cards separated by newlines. Must be parseable by astropy.")
args = parser.parse_args()
if not (2 <= len(args.dimensions) <= 4):
sys.exit("At least two dimensions required. At most 4 dimensions allowed.")
if not args.output:
args.output = "image-%s.fits" % "-".join(str(d) for d in args.dimensions)
make_image(args)