Merge pull request #2235 from quantified-uncertainty/complex-error-on…

…-power-fix

Support dist^integer for negative values

.changeset/tough-flowers-tap.md

@@ -0,0 +1,5 @@
+---
+"@quri/squiggle-lang": patch
+---
+
+Support pow(dist, integer) on distributions with negative samples

packages/squiggle-lang/__tests__/library/distOperations_test.ts

@@ -0,0 +1,14 @@
+import { testEvalToBe } from "../helpers/reducerHelpers.js";
+
+// TODO - it'd be useful to do most of `./sym_test.ts` tests here, without `Sym.` prefix
+describe("Dist operations", () => {
+  describe("Power on negative samples", () => {
+    // ok when power is integer
+    testEvalToBe("normal(-100, 1) ^ 2", "Sample Set Distribution");
+    // fails when power is not an integer
+    testEvalToBe(
+      "normal(-100, 1) ^ 2.5",
+      "Error(Distribution Math Error: Operation returned complex result)"
+    );
+  });
+});

packages/squiggle-lang/src/operation.ts

@@ -52,11 +52,14 @@ export const Convolution = {
 type OperationFn = (a: number, b: number) => result<number, OperationError>;
 
 const power: OperationFn = (a, b) => {
-  if (a >= 0) {
-    return Ok(a ** b);
-  } else {
+  const result = a ** b;
+  if (Number.isNaN(result)) {
+    if (Number.isNaN(a) || Number.isNaN(b)) {
+      return Ok(result); // bad, but the issue is upstream of `power` operation
+    }
     return Result.Err(new ComplexNumberError());
   }
+  return Ok(result);
 };
 
 const add: OperationFn = (a, b) => Ok(a + b);