How to use Python functions in Mathematica functions in Python notebooks?

Question

I want to plot things with Mathematica, inside an ipynb Python notebook. I can easily plot functions that are expressed in the Wolfram Language, e.g. with

from wolframclient.evaluation import WolframLanguageSession
from wolframclient.language import wlexpr, wl
from PIL import Image
import io

with WolframLanguageSession() as session:
    plot = session.evaluate(wlexpr('Plot[Sin[x], {x,0, Pi}]'))
    img_data = session.evaluate(wl.ExportByteArray(plot, 'PNG'))
    img = Image.open(io.BytesIO(img_data))
img

but what I really want to do is to plot certain outputs of my machine learning model. Let's say I wrap that into some python function like

def compute_something(x: float, y: float) -> float:
    # ...

and now I want to create a ContourPlot of this function. Something like this:

session.evaluate(wlexpr(
    'ContourPlot[computeSomething[x, y], {x, 0, 1}, {y, 0, 1}]'
))

But how do I pass the numeric python function compute_something to the session so it is represented as computeSomething in Mathematica? Remember that I am in an iPython notebook with some Python kernel already running and holding this function.

(P.S.: I really really don't want to use matplotlib.)

Answer 1

Assuming we have a sample function in python:

def computeSomething(x,y):
    return x+y

Here are three possible solutions:

1 - Easiest

Generate data points in Python, export the table, and use ListContourPlot (assuming you have numpy):

In Python

import numpy as np

# generate a table
data=[[x,y,computeSomething(x,y)] for x in np.arange(0,1,.2) for y in np.arange(0,1,.25)]

# export the data
np.savetxt(r'C:\file.csv',d, fmt='%s', delimiter=",")

In Mathematica

data = Import["C:\\file.csv", "CSV"];

ListContourPlot[data]

Output:

2 - Advanced (blocking)

Use a socket to communicate results from Python to Wolfram. We'll use 0MQ instead of the plain python TCP socket. You should install pyzmq (the same library Mathematica uses for Python ExternalEvaluate)

In Python

Now, we'll create a socket and bind your function, so on receiving input, output would be calculated and the result would be sent back.

import zmq

HOST = "127.0.0.1"  # localhost
PORT = 65435  # non-privileged ports are > 1023

context = zmq.Context()
socket = context.socket(zmq.REP)
socket.bind(f"tcp://{HOST}:{PORT}")

with socket:
    while True:
        
        message = socket.recv().decode()
        if message=='quit':
            break

        # print (f"[LOG] Received request: \"{message}\"")

        x,y=message.split(' ')
        x=float(x)
        y=float(y)
        socket.send_string(str(computeSomething(x,y)))

In Mathematica

We'll connect to the socket and define a connector function:

socket = SocketConnect["tcp://127.0.0.1:65435", {"ZMQ", "Request"}]

ClearAll[computeSomething];

computeSomething[x_?NumberQ, y_?NumberQ] := (WriteString[socket, ToString[x] <> " " <> ToString[y]];
  Internal`StringToMReal[ReadString[socket]])

Plot:

ContourPlot[computeSomething[x, y], {x, 0, 1}, {y, 0, 1}, 
  MaxRecursion -> 0, PlotPoints -> 2, 
  PerformanceGoal -> "Speed"]

Output:

The Python section is running and is waiting for input, so If you're done with your Python side, send "quit" to stop that.

WriteString[socket, "quit"]

And close Mathematica socket:

Close[socket];

3 - Advanced (non-blocking)

With my limited knowledge, here is another approach that uses IPython IOLoop.

In Python

import zmq
from zmq.eventloop.future import Context
from zmq.eventloop.zmqstream import ZMQStream

HOST = "127.0.0.1"  # localhost
PORT = 65435  # non-privileged ports are > 1023

context = Context.instance()
socket = context.socket(zmq.REP)
socket.bind(f"tcp://{HOST}:{PORT}")
stream = ZMQStream(socket)

def temp(message):
    x,y=message.result()[0].decode().split(' ')
    x=float(x)
    y=float(y)
    stream.send_string(str(computeSomething(x,y)))

stream.on_recv(temp)

# ioloop.IOLoop.current().start() is not needed because IPython is already started a loop, so it'll raise error

Use Solution 2 Mathematica code to connect to the server and plot the result.

Notes:

• I used minimal quality in the Advance section and it took on average under 1 second (sometimes the output was not right, re-evaluating will help)
• On higher quality, it'll take a long time to calculate (you can mix the above solutions to create the table inside Mathematica)
• Internal StringToMReal was used to convert String to Real, it's undocumented and is not guaranteed to work on all versions (in fact it was renamed for consistency)