34 lines
No EOL
988 B
Python
34 lines
No EOL
988 B
Python
from .utils_functions import *
|
|
|
|
def prepare_data(batch,
|
|
scale_phi_ini:float = 1000.0,
|
|
scale_delta_ini:float = 12.0,
|
|
scale_target:float = 50.0,
|
|
lin_threshold_target:float = 2.,
|
|
):
|
|
|
|
target_transform = CustomSymlogTransform(threshold=lin_threshold_target)
|
|
|
|
# delta_ini = batch['input'][:, [0], :, :, :]
|
|
phi_ini = batch['input'][:, [1], :, :, :]
|
|
D1 = batch['style'][:, [0], None, None, None]
|
|
# D2 = batch['style'][:, [1], None, None, None]
|
|
gravpot = batch['target'][:, [0], :, :, :]
|
|
|
|
|
|
_input = batch['input']
|
|
_input[:, 0, :, :, :] /= scale_delta_ini
|
|
_input[:, 1, :, :, :] /= scale_phi_ini
|
|
|
|
_target = (gravpot/D1 - phi_ini)/D1
|
|
_target /= scale_target
|
|
_target = target_transform(_target)
|
|
|
|
_style = batch['style']
|
|
|
|
return {
|
|
'input': _input,
|
|
'target': _target,
|
|
'style': _style
|
|
}
|
|
|