feat: mathtransform with jscodeshift

[strout18]

Description

The code in this branch is meant to allow users to write constraints/objs in basic math, and translate them afterward to the specific optimization framework. See the README in react-renderer/src/mathtransform for more details.

Ideally it would run each time a file is changed (or at least a file from a specific list of files we want to transform), so hot-reloading functionality. It’s currently built to run on the command line (sample line: jscodeshift bugtest.ts -t toCustomAD.ts -p -d) , but I think it could also be modified to use the API functions (more detail is here). Another thing to note (also mentioned in the README) is that jscodeshift modifies in place the file you call it on (unless you call a dry run). So if we wanted to keep an unmodified version of that file we would need to make a copy. The CLI for jscodeshift apparently does not support redirecting the output (though the API does) But the program doesn’t change much code, and which code it changes is well-defined and easily reversible so maybe storing the old version is not necessary.

Sample Input/Output

Sample Input

    export const objDict = {
        // mathtrans
        equal: (x: number, y: number) => squared(x - y),

        // mathtrans
        above: ([t1, top]: [string, any], [t2, bottom]: [string, any], offset = 100) =>
            // (getY top - getY bottom - offset) ^ 2
            squared(top.y.contents - bottom.y.contents - varOf(offset))
    }

Sample Output

    export const objDict = {
        // mathtrans
        equal: (x: VarAD, y: VarAD) => squared(sub(x, y)),

        // mathtrans
        above: ([t1, top]: [string, any], [t2, bottom]: [string, any], offset = 100) =>
            // (getY top - getY bottom - offset) ^ 2
            squared(sub(sub(top.y.contents, bottom.y.contents), varOf(offset)))
    }

Remaining issues

Dependencies

I’m not sure what to do with the new dependencies. I was just originally working on this project outside of the Penrose directory, so I had to download the jscodeshift package and make my own tsconfig file to get it to work. I assume it should be merged into the main dependency list and node_modules folder but I'm not sure how to do that.

[strout18]

@wodeni

[wodeni]

The transformer looks good to me. I mainly want to know more about how the comment marker works and discuss a bit more about integration with the main system.

On the latter, one thought is to use ttsc (typescript compiler wrapper that allows plugins), or wait until tsc supports plugins and custom transforms. I also found an article about integrating custom transforms, which uses ttsc.

I still don't fully understand how typescript's built-in transforms are supported and if jscodeshift can work with ttsc. Let me know if you experimented with any of the above.

Anyway, really cool work ❤️ !

[wodeni]

Much of this module remains untyped. There seems to be @types/jscodeshift on npm. Do you think it'll help with readability/error prevention?

[strout18]

That's a great idea, I definitely think it will.

[wodeni]

Would be great to clarify the usage of this marker. Looks like it will transform any functions following this marker? Can you transform only some lines in a function then?
I wonder if decorators can be a viable alternative, but they are probably not expressive enough for meta tasks like this.
Finally, something like differentiable or autodiff would be more meaningful markers here?

[strout18]

It will actually transform the nearest AST node below the marker. If that's a function, or a statement, or something else. It checks the LeadingComments node that the parser attaches to another AST node, which I believe is whatever the main type of a line is (method call, function, variable declaration, etc.).

Sample Input:

export const objDict = {
    above: (x: number[]): number => {
      //autodiff
      x = t2 - z * 34 + (2 / 4)
      y = a + bigint
    }
}

Sample Output:

export const objDict = {
    above: (x: number[]): number => {
      //autodiff
      x = add(sub(t2, mul(z, 34)), div(2, 4))
      y = a + bigint
    }
}

Renamed in 635db53

[wodeni]

Maybe link to an API reference for future use. Seems like jscodeshift is using Mozilla's parser API?

[wodeni]

Maybe not a ton better: at least make sure getTransformData returns Node | undefined, which will let the type checker do some checking. Also can you destructure transformData to give the outputs meaningful names, e.g. const [operator, transformer] = getTransformData(node)?

[kai-qu]

Thanks for doing this Stella! Super exciting work and it will be really helpful for our users; I definitely want to integrate it into our system as soon as we can 🙂 I was able to repro the example too.

Some fixes that would make life easier for us:

• Very commonly we want to transform JS patterns into custom AD patterns, not just single operators. So, is it possible to implement nodes that transform patterns, not just single ops, and enable the user to specify this pattern in mathtransform? Here are some examples:

Math.pow(x, y) => pow(x, y)
Math.pow(x, 2) => squared(x)
_.sum(xs) => addN(x)
cond ? a : b => ifCond(cond, a, b)

Maybe @wodeni can comment on how this might be done?

• Anything that's a number actually needs varOf auto-applied to it (to change the type from number to VarAD), and there should be no varOf in the input code. Would you be able to do that? Here I've changed your sample code to reflect both this and the previous ask:

Sample input:

export const objDict = {
    // mathtrans
    equal: (x: number, y: number) => Math.pow(x - y, 2), // continue to have no varOf here; also changed to use JS function

    // mathtrans
    above: ([t1, top]: [string, any], [t2, bottom]: [string, any], offset = 100) =>
        // (getY top - getY bottom - offset) ^ 2
        Math.pow(top.y.contents - bottom.y.contents - offset, 2) // <-- changed: no varOf on offset; also changed to use JS function
}

Sample output:

    export const objDict = {
        // mathtrans
        equal: (x: VarAD, y: VarAD) => squared(sub(varOf(x), varOf(y))), // <-- changed: need varOf here

        // mathtrans
        above: ([t1, top]: [string, any], [t2, bottom]: [string, any], offset = 100) =>
            // (getY top - getY bottom - offset) ^ 2
            squared(sub(sub(top.y.contents, bottom.y.contents), varOf(offset))) // <-- unchanged: continue to have varOf here
    }

If wrapping things in varOf is too difficult, I may be able to write it as a wrapper around my functions, but doing it here would be more modularized.

On documentation:

• Can you note that the user needs to install jscodeshift?
• The more detail is here link seems to be broken
• You can also note that nesting the // mathtransform annotation works (on nested AST nodes), they don't have to be top-level function definitions

For us (@wodeni @maxkrieger and I)

• We need to figure out how and when to run this code. Ideally it runs on all user-defined TS functions in the Style programs that are inputted (including via imports), as well as all the Penrose TS contrib files. What do y'all think?

[strout18]

Thanks for the sample input and output, btw, those were super helpful!

Very commonly we want to transform JS patterns into custom AD patterns, not just single operators

This can totally be done. I'll start working on it. The transforms have to be written differently for each type of expression though, so any examples of any other types of expressions (beyond the method call type and ternary expression type you provided below) that you think should be translatable would be great. For one, I would think it's also useful to be able to rename a method, like norm() to ADNorm(), so I'll add that to the list.

Math.pow(x, y) => pow(x, y)
Math.pow(x, 2) => squared(x)
_.sum(xs) => addN(x)
cond ? a : b => ifCond(cond, a, b) 
norm(x) => ADNorm(x)

Anything that's a number actually needs varOf auto-applied to it (to change the type from number to VarAD)

This could be trickier. I could set it up so it would work for that example and for a lot of situations that are similar. But JSCodeshift doesn't know about the inferred types of variables, only the actual (untyped) AST node of the variable itself, so if the "number" nodes look exactly like nodes of another type, it would transform the other undesirable node as well. For instance, b + a would get transformed into varOf(b) + varOf(a) even if you didn't want to transform a. So I guess my point is, do you think that kind of scenario is possible within the scope of transformation (so within a function or wherever else it might be required)? If that's likely then, yeah, this may be too difficult to implement.

Also isvarOf applied to both variables and numbers? So 1 + b is translated to varOf(1) + varOf(b)?

[kai-qu]

The transforms have to be written differently for each type of expression though, so any examples of any other types of expressions (beyond the method call type and ternary expression type you provided below) that you think should be translatable would be great.

Awesome, thanks, this will be super useful! Nothing else comes to mind, except that the translation rules might need to apply with some specificity/precedence if more than one of them matches, e.g.

Math.pow(x, y) => pow(x, y)
Math.pow(x, 2) => squared(x)

Both would match Math.pow(v,2), but the latter is a better fit. You could just ask users to put the more "specific" rules later in the list, and always pick the latest rule that applies to a node.

For one, I would think it's also useful to be able to rename a method, like norm() to ADNorm(), so I'll add that to the list.

Agreed.

But JSCodeshift doesn't know about the inferred types of variables, only the actual (untyped) AST node of the variable itself, so if the "number" nodes look exactly like nodes of another type, it would transform the other undesirable node as well. ... So I guess my point is, do you think that kind of scenario is possible within the scope of transformation (so within a function or wherever else it might be required)?

Ok, I'm gonna backtrack on this one. On second thought, I would prefer mathtransform to not convert any variables or constants to AD vars, because either they are already going to be AD vars in the system, or the AD function used can inspect the type of the input at runtime and decide what wrapper to use (if any).

What should change is the name of the function used; it will probably be something like addWrapper instead of add (but I can make those changes when I implement this).

So, in the end, it's just fine to convert something like Math.pow(x, 3) to powWrapper(x, 3), NOT pow(varOf(x), constOf(3)).

BTW, I forgot to confirm: you are converting function type signatures too, right? So anything that has a number type gets a VarAD type? (That's the right thing to do, and if your code already does that, you can keep it that way :-)) Example:

equal: (x: number, y: number, offset: number, [t2, bot]: [string, any])
=>
equal: (x: VarAD, y: VarAD, offset: VarAD, [t2, bot]: [string, any])

Let me know if anything is confusing, or when this is ready to re-review!

[strout18]

@hypotext @wodeni This is ready for re-review! Sorry this took so long to get done. In order to implement those changes you requested, I had to do a fairly significant refactoring. But it should be much easier to adapt to your needs now. I also added to the README a section breaking down the code further. The current functionality is also listed in the README. It should be pretty easy to change it to what exactly you want (e.g. add to addWrapper).

[kai-qu]

Thanks Stella, this looks very helpful! I did a quick check of the functionality--can you just make sure that the pattern replacement works on nested nodes? For example, replacing Math.pow doesn't seem to work when it's nested:

Input code: (I know the # args to vdistsq is wrong but shouldn't matter for the purposes of this test)

    // autodiff
    testpow3: ([t1, s1]: [string, any], [t2, s2]: [string, any]) =>
        ops.vdistsq(Math.pow(x - 3, y)),

    // autodiff
    testpow4: ([t1, s1]: [string, any], [t2, s2]: [string, any]) =>
        ops.vdistsq(Math.pow(x, y - 3)),

Output code:

    // autodiff
    testpow3: ([t1, s1]: [string, any], [t2, s2]: [string, any]) =>
        ops.vdistsq(Math.pow(sub(x, 3), y)), // only replaced with `sub`

    // autodiff
    testpow4: ([t1, s1]: [string, any], [t2, s2]: [string, any]) =>
        ops.vdistsq(Math.pow(x, sub(y, 3))),

Expected code:

    // autodiff
    testpow3: ([t1, s1]: [string, any], [t2, s2]: [string, any]) =>
        ops.vdistsq(pow(sub(x, 3), y)), // should replace `pow` as well

    // autodiff
    testpow4: ([t1, s1]: [string, any], [t2, s2]: [string, any]) =>
        ops.vdistsq(pow(x, sub(y, 3))),

I'd just ask that you test for this kind of case more thoroughly, for patterns nested in other patterns as well (e.g. a pow in a cond in a plus in a vdistsq... etc.)

[kai-qu]

Btw, is there some way to turn off the auto-build for the test typescript files for this feature @wodeni?

I tried to push my test cases but I get a bunch of build errors on files like src/mathtransform/bugtest.ts, which should basically just be treated as text, and not as compilable ts files with all references defined, etc.

[wodeni]

Btw, is there some way to turn off the auto-build for the test typescript files for this feature @wodeni?

I tried to push my test cases but I get a bunch of build errors on files like src/mathtransform/bugtest.ts, which should basically just be treated as text, and not as compilable ts files with all references defined, etc.

There's // @ts-nocheck or you can explicitly ignore it in tsconfig.json under the exclude key.

If they are test cases, maybe it's better to include them in our test suite eventually btw.

[kai-qu]

Thanks. @strout18 Can you try and add the tests and commit them?

I tried various things (including decorating the files mathtransform/Constraints.ts, mathtransform/bugtest.ts, and mathtransform/mathtest.ts with // @ts-nocheck, as well as adding those files to an exclude list in both mathtransform/tsconfig.json and ROOT/tsconfig.json) without success; I'm not able to commit my changes.

[maxkrieger]

You can use the --no-verify flag in git iirc

…

On Dec 9, 2020, at 22:14, Katherine Ye ***@***.***> wrote: Thanks. @strout18 <https://github.com/strout18> Can you try and add the tests and commit them? I tried various things (including decorating the files mathtransform/Constraints.ts, mathtransform/bugtest.ts, and mathtransform/mathtest.ts with // @ts-nocheck, as well as adding those files to an exclude list in both mathtransform/tsconfig.json and ROOT/tsconfig.json) without success; I'm not able to commit my changes. — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub <#389 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAUJSGVEO2A6PL6UPD23U2DSUA4J5ANCNFSM4SS32ZKQ>.

[strout18]

@wodeni @hypotext This is ready for re-review. I added a bunch of tests in the tests.ts file and they all pas.

[kai-qu]

Hey Stella, I had a look at this and it looks great!

One final request: Can you rewrite your tests (testkey.ts) in the form of a Jest test suite? Basically, running the test suite should automatically check that each output of running the command jscodeshift tests.ts -t toCustomAD.ts -p -d matches the corresponding output in testkey.ts.

We've been adding tests for all the new parts of penrose-web, and it would make it much easier to check your code against future changes, instead of going through and manually verifying that the test output matches the expected testkey.ts result.

You can see an example test suite here: https://github.com/penrose/penrose/blob/style-compiler-port/penrose-web/src/compiler/Style.test.ts

It should be pretty quick to port testkey.ts to Jest, and then I'd just ask you verify that it integrates with the other Penrose tests, i.e. it runs when you do npm test. Then we can finally merge this!

(@wodeni Does Stella need to do anything special on her branch to get tests working with penrose-web, since this branch is before all the Style compiler port stuff?)

[wodeni]

We never ejected penrose-web, so jest should work just fine. If you run into problems, just let me know @strout18

[strout18]

@hypotext @wodeni

Okay, I set it up with Jest, so npm test works (from the __tests__ directory). Turns out JSCodeshift provides built-in utilities to test with Jest, so it only took two extra lines to get it working (although I did have to add a __testfixtures__ directory under penrose-web/src, so that's where that came from if anyone notices it). Loose ends to tie up:

• NPM stuff. When I started this project I didn't really know how to use Node.js, so I installed all the packages I needed directly into the folder I was working in and created my own tsconfig and package.json instead of adding the dependencies to Penrose overall. Still not quite sure what I need to do to get it all merged with the main project configuration.
• Parentheses: For some reason, after rearranging the code, JSCodeshift sometimes adds extra sets of parentheses around AST elements (it seems to be only inside binary operations), so instead of transforming 3 + 7 * 2 into add(3, mul(7, 2)) it returns add(3, (mul(7,2))). This obviously isn't a problem syntactically, as the meaning is preserved, but it's just a little confusing aesthetically. If we think it's necessary I can try to prevent that, but currently that seems like a lot of unnecessary work, as I'd have to go digging around in the JSCodeshift code since I'm pretty sure the problem is in the API and not on my end.

[kai-qu]

Hey Stella, thanks for doing this! I was able to replicate it but had to run npm install jscodeshift in the test directory first. Maybe good to note this in the README if the dependencies are not resolved.

NPM stuff. When I started this project I didn't really know how to use Node.js, so I installed all the packages I needed directly into the folder I was working in and created my own tsconfig and package.json instead of adding the dependencies to Penrose overall. Still not quite sure what I need to do to get it all merged with the main project configuration.

@wodeni Can you weigh in on how Stella should integrate the dependencies?

(After we get @wodeni's feedback it should be good to merge)

[wodeni]

Thanks for the update @strout18! The test suite worked for me as well after I manually install jscodeshift. See my comments for some lower-level questions. Hopefully after you address them, CI will pass again.

Re integration with the core: my understanding is one will need to run jscodeshift as a part of our build step eventually. We are finishing up the port to TS on #419. It's ok to merge this as an independent package for now, and I'll integrate it into the new system after #446.

[wodeni]

Looks great! @hypotext is there any comment from your end? I think it's good to merge.

[codecov]

Codecov Report

Merging #389 (ddbb8e4) into main (2c143ac) will increase coverage by 0.45%.
The diff coverage is 87.75%.

@@            Coverage Diff             @@
##             main     #389      +/-   ##
==========================================
+ Coverage   70.61%   71.06%   +0.45%     
==========================================
  Files          31       32       +1     
  Lines        5373     5520     +147     
  Branches     1029     1052      +23     
==========================================
+ Hits         3794     3923     +129     
- Misses       1574     1592      +18     
  Partials        5        5              

Impacted Files	Coverage Δ
packages/core/src/utils/toCustomAD.ts	87.75% <87.75%> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 2c143ac...ddbb8e4. Read the comment docs.

[kai-qu]

LGTM, feel free to merge! (It looks like there are minor package.json, yarn.lock etc. conflicts with main. Maybe @wodeni can weigh in on how to resolve those?)

[kai-qu]

Checking in on this -- could we get this PR merged before any more merge conflicts arise? Looks like it's just some package dependency files that have conflicts. Maybe we can resolve during today's "office" time? @wodeni @maxkrieger

[wodeni]

I attempted to merge this branch earlier, but ran into the following issue:

The test suite worked out just fine before when the new files ran with penrose-web, which had CRA installed. Now that @penrose/core is an independent package without CRA, we lost all the mysterious jest support that came with CRA. The issue now is that import { defineTest } from "jscodeshift/dist/testUtils" no longer works. It either doesn't typecheck possible due to the jscodeshift/dist/testUtils (jest doesn't really handle nested import by itself and we don't know how CRA did it), or the test suite fails without any output if I disable typechecking in mathtransform.test.ts (see CI output).

Turns out that the default export of toCustomAD was the problem. Fixed.