split up elementwise atoms (#57)

Author: dance858

include/elementwise_full_dom.h

@@ -0,0 +1,34 @@
+#ifndef ELEMENTWISE_FULL_DOM_H
+#define ELEMENTWISE_FULL_DOM_H
+
+#include "expr.h"
+
+/* Helper function to initialize an elementwise expr
+ * (can be used with derived types) */
+void init_elementwise(expr *node, expr *child);
+
+expr *new_exp(expr *child);
+expr *new_sin(expr *child);
+expr *new_cos(expr *child);
+expr *new_sinh(expr *child);
+expr *new_tanh(expr *child);
+expr *new_asinh(expr *child);
+expr *new_logistic(expr *child);
+expr *new_power(expr *child, double p);
+expr *new_xexp(expr *child);
+expr *new_normal_cdf(expr *child);
+
+/* the jacobian and wsum_hess for elementwise full domain
+   atoms are always initialized in the same way and
+   implement the chain rule in the same way */
+void jacobian_init_elementwise(expr *node);
+void eval_jacobian_elementwise(expr *node);
+void wsum_hess_init_elementwise(expr *node);
+void eval_wsum_hess_elementwise(expr *node, const double *w);
+expr *new_elementwise(expr *child);
+
+/* no elementwise atoms are affine according to our
+   convention, so we can have a common implementation */
+bool is_affine_elementwise(const expr *node);
+
+#endif /* ELEMENTWISE_FULL_DOM_H */

include/elementwise_restricted_dom.h

@@ -0,0 +1,18 @@
+#ifndef ELEMENTWISE_RESTRICTED_DOM_H
+#define ELEMENTWISE_RESTRICTED_DOM_H
+
+#include "expr.h"
+
+/* Shared init functions for restricted domain atoms
+ * (variable-child only, no linear operator support) */
+void jacobian_init_restricted(expr *node);
+void wsum_hess_init_restricted(expr *node);
+bool is_affine_restricted(const expr *node);
+expr *new_restricted(expr *child);
+
+expr *new_log(expr *child);
+expr *new_entr(expr *child);
+expr *new_atanh(expr *child);
+expr *new_tan(expr *child);
+
+#endif /* ELEMENTWISE_RESTRICTED_DOM_H */

include/elementwise_univariate.h

@@ -1,22 +1,4 @@
-/*
- * Copyright 2026 Daniel Cederberg and William Zhang
- *
- * This file is part of the DNLP-differentiation-engine project.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-#include "elementwise_univariate.h"
-#include "expr.h"
+#include "elementwise_full_dom.h"
 #include "subexpr.h"
 #include <stdio.h>
 #include <stdlib.h>
@@ -121,10 +103,6 @@ bool is_affine_elementwise(const expr *node)
     return false;
 }
 
-/* Helper function to initialize an already-allocated expr for elementwise operations
- * This is called when a power_expr or other type-specific struct is allocated
- * and we need to initialize the base expr fields
- */
 void init_elementwise(expr *node, expr *child)
 {
     /* Initialize base fields */

src/elementwise_univariate/common.c src/elementwise_full_dom/common.csrc/elementwise_univariate/common.c renamed to src/elementwise_full_dom/common.c

@@ -1,21 +1,4 @@
-/*
- * Copyright 2026 Daniel Cederberg and William Zhang
- *
- * This file is part of the DNLP-differentiation-engine project.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-#include "elementwise_univariate.h"
+#include "elementwise_full_dom.h"
 #include <math.h>
 #include <string.h>
 

src/elementwise_univariate/exp.c src/elementwise_full_dom/exp.csrc/elementwise_univariate/exp.c renamed to src/elementwise_full_dom/exp.c

@@ -1,4 +1,4 @@
-#include "elementwise_univariate.h"
+#include "elementwise_full_dom.h"
 #include <assert.h>
 #include <math.h>
 
@@ -114,41 +114,3 @@ expr *new_asinh(expr *child)
     node->local_wsum_hess = asinh_local_wsum_hess;
     return node;
 }
-
-/* ----------------------- atanh ----------------------- */
-static void atanh_forward(expr *node, const double *u)
-{
-    node->left->forward(node->left, u);
-    for (int i = 0; i < node->size; i++)
-    {
-        node->value[i] = atanh(node->left->value[i]);
-    }
-}
-
-static void atanh_local_jacobian(expr *node, double *vals)
-{
-    expr *child = node->left;
-    for (int j = 0; j < node->size; j++)
-    {
-        vals[j] = 1.0 / (1.0 - child->value[j] * child->value[j]);
-    }
-}
-
-static void atanh_local_wsum_hess(expr *node, double *out, const double *w)
-{
-    double *x = node->left->value;
-    for (int j = 0; j < node->size; j++)
-    {
-        double c = 1.0 - x[j] * x[j];
-        out[j] = w[j] * (2.0 * x[j]) / (c * c);
-    }
-}
-
-expr *new_atanh(expr *child)
-{
-    expr *node = new_elementwise(child);
-    node->forward = atanh_forward;
-    node->local_jacobian = atanh_local_jacobian;
-    node->local_wsum_hess = atanh_local_wsum_hess;
-    return node;
-}

src/elementwise_univariate/hyperbolic.c src/elementwise_full_dom/hyperbolic.csrc/elementwise_univariate/hyperbolic.c renamed to src/elementwise_full_dom/hyperbolic.c

@@ -1,21 +1,4 @@
-/*
- * Copyright 2026 Daniel Cederberg and William Zhang
- *
- * This file is part of the DNLP-differentiation-engine project.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-#include "elementwise_univariate.h"
+#include "elementwise_full_dom.h"
 #include <math.h>
 
 static void forward(expr *node, const double *u)

src/elementwise_univariate/logistic.c src/elementwise_full_dom/logistic.csrc/elementwise_univariate/logistic.c renamed to src/elementwise_full_dom/logistic.c

@@ -1,21 +1,4 @@
-/*
- * Copyright 2026 Daniel Cederberg and William Zhang
- *
- * This file is part of the DNLP-differentiation-engine project.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-#include "elementwise_univariate.h"
+#include "elementwise_full_dom.h"
 #include <math.h>
 
 #ifndef M_PI
@@ -53,7 +36,6 @@ static void local_wsum_hess(expr *node, double *out, const double *w)
     double *x = node->left->value;
     for (int j = 0; j < node->size; j++)
     {
-        /* could avoid recomputing this (like in logistic) */
         double phi = INV_SQRT_2PI * exp(-0.5 * x[j] * x[j]);
         out[j] = w[j] * (-x[j] * phi);
     }

src/elementwise_univariate/normal_cdf.c src/elementwise_full_dom/normal_cdf.csrc/elementwise_univariate/normal_cdf.c renamed to src/elementwise_full_dom/normal_cdf.c

@@ -1,21 +1,4 @@
-/*
- * Copyright 2026 Daniel Cederberg and William Zhang
- *
- * This file is part of the DNLP-differentiation-engine project.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-#include "elementwise_univariate.h"
+#include "elementwise_full_dom.h"
 #include "subexpr.h"
 #include <math.h>
 #include <stdlib.h>

src/elementwise_univariate/power.c src/elementwise_full_dom/power.csrc/elementwise_univariate/power.c renamed to src/elementwise_full_dom/power.c

@@ -1,21 +1,4 @@
-/*
- * Copyright 2026 Daniel Cederberg and William Zhang
- *
- * This file is part of the DNLP-differentiation-engine project.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-#include "elementwise_univariate.h"
+#include "elementwise_full_dom.h"
 #include <math.h>
 
 /* ----------------------- sin ----------------------- */
@@ -89,43 +72,3 @@ expr *new_cos(expr *child)
     node->local_wsum_hess = cos_local_wsum_hess;
     return node;
 }
-
-/* ----------------------- tan ----------------------- */
-static void tan_forward(expr *node, const double *u)
-{
-    node->left->forward(node->left, u);
-    for (int i = 0; i < node->size; i++)
-    {
-        node->value[i] = tan(node->left->value[i]);
-    }
-}
-
-static void tan_local_jacobian(expr *node, double *vals)
-{
-    expr *child = node->left;
-    for (int j = 0; j < node->size; j++)
-    {
-        double c = cos(child->value[j]);
-        vals[j] = 1.0 / (c * c);
-    }
-}
-
-static void tan_local_wsum_hess(expr *node, double *out, const double *w)
-{
-    double *x = node->left->value;
-
-    for (int j = 0; j < node->size; j++)
-    {
-        double c = cos(x[j]);
-        out[j] = 2.0 * w[j] * node->value[j] / (c * c);
-    }
-}
-
-expr *new_tan(expr *child)
-{
-    expr *node = new_elementwise(child);
-    node->forward = tan_forward;
-    node->local_jacobian = tan_local_jacobian;
-    node->local_wsum_hess = tan_local_wsum_hess;
-    return node;
-}