File size: 1,950 Bytes
b3d42fc |
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 |
from tensorflow.keras.models import load_model
import tensorflow as tf
from tensorflow.keras.saving import register_keras_serializable
from tensorflow.keras import layers, models, backend as K
@register_keras_serializable()
class SSTAmplifier(tf.keras.layers.Layer):
def __init__(self, threshold=28.0, scale=0.1, **kwargs):
super().__init__(**kwargs)
self.threshold = threshold
self.scale = scale
def call(self, inputs):
sst = inputs[:, 0]
factor = tf.sigmoid((sst - self.threshold) * self.scale)
mod = 1.0 + 0.3 * factor
return tf.expand_dims(mod, -1)
@register_keras_serializable()
class ShearSuppressor(tf.keras.layers.Layer):
def __init__(self, threshold=14.0, scale=0.2, **kwargs):
super().__init__(**kwargs)
self.threshold = threshold
self.scale = scale
def call(self, inputs):
shear = inputs[:, 3]
suppress = tf.sigmoid((self.threshold - shear) * self.scale)
mod = 1.0 - 0.25 * suppress
return tf.expand_dims(mod, -1)
@register_keras_serializable()
class VorticityActivator(tf.keras.layers.Layer):
def __init__(self, threshold=1.2, scale=1.0, **kwargs):
super().__init__(**kwargs)
self.threshold = threshold
self.scale = scale
def call(self, inputs):
vort = inputs[:, 4]
activate = tf.sigmoid((vort - self.threshold) * self.scale)
mod = 1.0 + 0.2 * activate
return tf.expand_dims(mod, -1)
@register_keras_serializable()
class ModulationMixer(tf.keras.layers.Layer):
def call(self, inputs):
sst_mod, shear_mod, vort_mod = inputs
product = sst_mod * shear_mod * vort_mod
smooth = 1.0 + 0.25 * tf.tanh(product - 1.0)
return smooth
CUSTOM_OBJECTS = {
'ModulationMixer': ModulationMixer,
'VorticityActivator': VorticityActivator,
'ShearSuppressor': ShearSuppressor,
'SSTAmplifier': SSTAmplifier
} |