| Geant4 Cross Reference |
1 """
2 ** generate **
3 generate showers using a saved VAE model
4 """
5 import argparse
6
7 import numpy as np
8 import tensorflow as tf
9 from tensorflow.python.data import Dataset
10
11 from core.constants import GLOBAL_CHECKPOINT_D
12 from utils.gpu_limiter import GPULimiter
13 from utils.preprocess import get_condition_arr
14
15
16 def parse_args():
17 argument_parser = argparse.ArgumentParser(
18 argument_parser.add_argument("--geometry",
19 argument_parser.add_argument("--energy", t
20 argument_parser.add_argument("--angle", ty
21 argument_parser.add_argument("--events", t
22 argument_parser.add_argument("--epoch", ty
23 argument_parser.add_argument("--study-name
24 argument_parser.add_argument("--max-gpu-me
25 argument_parser.add_argument("--gpu-ids",
26 args = argument_parser.parse_args()
27 return args
28
29
30 # main function
31 def main():
32 # 0. Parse arguments.
33 args = parse_args()
34 energy = args.energy
35 angle = args.angle
36 geometry = args.geometry
37 events = args.events
38 epoch = args.epoch
39 study_name = args.study_name
40 max_gpu_memory_allocation = args.max_gpu_m
41 gpu_ids = args.gpu_ids
42
43 # 1. Set GPU memory limits.
44 GPULimiter(_gpu_ids=gpu_ids, _max_gpu_memo
45
46 # 2. Load a saved model.
47
48 # Create a handler and build model.
49 # This import must be local because otherw
50 from core.model import VAEHandler
51 vae = VAEHandler()
52
53 # Load the saved weights
54 weights_dir = f"VAE_epoch_{epoch:03}" if e
55 vae.model.load_weights(f"{GLOBAL_CHECKPOIN
56
57 # The generator is defined as the decoder
58 generator = vae.model.decoder
59
60 # 3. Prepare data. Get condition values. S
61 # latent space dimension). Gather them int
62 # on the Dataset objects. Disable AutoShar
63 e_cond, angle_cond, geo_cond = get_conditi
64
65 z_r = np.random.normal(loc=0, scale=1, siz
66
67 data = ((z_r, e_cond, angle_cond, geo_cond
68
69 data = Dataset.from_tensor_slices(data)
70
71 batch_size = BATCH_SIZE_PER_REPLICA
72
73 data = data.batch(batch_size)
74
75 options = tf.data.Options()
76 options.experimental_distribute.auto_shard
77 data = data.with_options(options)
78
79 # 4. Generate showers using the VAE model.
80 generated_events = generator.predict(data)
81
82 # 5. Save the generated showers.
83 np.save(f"{GEN_DIR}/VAE_Generated_Geo_{geo
84
85
86 if __name__ == "__main__":
87 exit(main())