diff --git a/packages/squiggle-lang/src/rescript/FunctionRegistry/Library/FR_Pointset.res b/packages/squiggle-lang/src/rescript/FunctionRegistry/Library/FR_Pointset.res index 4f4e1731..c0270e48 100644 --- a/packages/squiggle-lang/src/rescript/FunctionRegistry/Library/FR_Pointset.res +++ b/packages/squiggle-lang/src/rescript/FunctionRegistry/Library/FR_Pointset.res @@ -23,6 +23,31 @@ let inputsTodist = (inputs: array, makeDist) => { expressionValue } +module Internal = { + type t = PointSetDist.t + + let toType = (r): result< + ReducerInterface_InternalExpressionValue.t, + Reducer_ErrorValue.errorValue, + > => + switch r { + | Ok(r) => Ok(Wrappers.evDistribution(PointSet(r))) + | Error(err) => Error(REOperationError(err)) + } + + let doLambdaCall = (aLambdaValue, list, environment, reducer) => + switch Reducer_Expression_Lambda.doLambdaCall(aLambdaValue, list, environment, reducer) { + | Ok(IEvNumber(f)) => Ok(f) + | _ => Error(Operation.SampleMapNeedsNtoNFunction) + } + + let mapY = (pointSetDist: t, aLambdaValue, env, reducer) => { + let fn = r => doLambdaCall(aLambdaValue, list{IEvNumber(r)}, env, reducer) + let foo = PointSetDist.T.mapYResult(~fn, pointSetDist) + foo->toType + } +} + let library = [ Function.make( ~name="fromDist", @@ -53,6 +78,26 @@ let library = [ ], (), ), + Function.make( + ~name="mapY", + ~nameSpace, + ~requiresNamespace=true, + ~examples=[`PointSet.mapY(mx(normal(5,2)), {|x| x + 1})`], + ~output=ReducerInterface_InternalExpressionValue.EvtDistribution, + ~definitions=[ + FnDefinition.make( + ~name="mapY", + ~inputs=[FRTypeDist, FRTypeLambda], + ~run=(inputs, _, env, reducer) => + switch inputs { + | [IEvDistribution(PointSet(dist)), IEvLambda(lambda)] => Internal.mapY(dist, lambda, env, reducer)->E.R2.errMap(_ => "") + | _ => Error(impossibleError) + }, + (), + ), + ], + (), + ), Function.make( ~name="makeContinuous", ~nameSpace,