Support layer-wise locally adaptive activation function

Author: lululxvi

deepxde/maps/activations.py

@@ -2,17 +2,37 @@
 from __future__ import division
 from __future__ import print_function
 
+from .. import config
 from ..backend import tf
 
 
 def linear(x):
     return x
 
 
+def layer_wise_locally_adaptive(activation, n=1):
+    """Layer-wise locally adaptive activation functions (L-LAAF).
+    Jagtap et al., arXiv preprint arXiv:1909.12228, 2019.
+
+    Examples:
+    To define a L-LAAF ReLU with the scaling factor ``n = 10``
+    ```python
+        n = 10
+        activation = f"LAAF-{n} relu"  # "LAAF-10 relu"
+    ```
+    """
+    a = tf.Variable(1 / n, dtype=config.real(tf))
+    return lambda x: activation(n * a * x)
+
+
 def get(identifier):
     if identifier is None:
         return linear
     if isinstance(identifier, str):
+        if identifier.startswith("LAAF"):
+            identifier = identifier.split()
+            n = float(identifier[0].split("-")[1])
+            return layer_wise_locally_adaptive(get(identifier[1]), n=n)
         return {
             "elu": tf.nn.elu,
             "relu": tf.nn.relu,