How to use NumPy slices in 2d for PyPlot's GridSpec

I am trying to create a plot with multiple subplots in PyPlot. For that I am using matplotlib’s GridSpec .
Some of the usage is given by NumPy slices as seen in their example from the docs:

fig8 = plt.figure(constrained_layout=False)
gs1 = fig8.add_gridspec(nrows=3, ncols=3, left=0.05, right=0.48, wspace=0.05)
f8_ax1 = fig8.add_subplot(gs1[:-1, :])
f8_ax2 = fig8.add_subplot(gs1[-1, :-1])
f8_ax3 = fig8.add_subplot(gs1[-1, -1])

It was shown here that slicing can be done by

f8_ax1 = fig8.add_subplot(get(gs1, pycall(pybuiltin("slice"), PyObject, 0,3)))

The above line is equivalent to gs1[0, 0:3].

But my question is, how can I do it in 2d? I mean how can I slice for example gs1[0:2, 0:3]?

Use a combination of get and pybuiltin("slice")

using PyPlot
using PyCall

fig = figure()
gs = fig.add_gridspec(2,2)

element(i,j) = get(gs, (i,j))
slice(i,j) = pycall(pybuiltin("slice"), PyObject, i,j)

ax = fig.add_subplot(element(0,0))
ax.plot(sin.(-pi:0.1:pi))

ax = fig.add_subplot(element(0,1))
ax.plot(cos.(-pi:0.1:pi))

ax = fig.add_subplot(element(1,slice(0,2)))
ax.plot(-pi:0.1:pi |> collect)

2 Likes

A convenient wrapper:

import PyPlot as plt; using PyPlot: @L_str, @py_str, @pyinclude, pyimport, pycall, PyDict, pybuiltin, PyObject

gridspec = pyimport("matplotlib.gridspec")

slice(args...) = pycall(pybuiltin("slice"), PyObject, args...)

# Wrapper struct for matplotlib GridSpec with Julia-style indexing
struct JulianGridSpec
   gs::PyObject  # Python GridSpec object
end

# Enable Julia-style indexing (1-based) on GridSpec objects
function Base.getindex(jgs::JulianGridSpec, i::Union{Int, UnitRange, Colon}, j::Union{Int, UnitRange, Colon})
   # Convert Julia indices to Python 0-based indices
   py_i = to_python_index(i)
   py_j = to_python_index(j)
   return get(jgs.gs, (py_i, py_j))
end

# Convert Julia integer to Python 0-based index
function to_python_index(idx::Int)
   return idx - 1
end

# Convert Julia range to Python slice
function to_python_index(idx::UnitRange)
   return slice(idx.start - 1, idx.stop)
end

# Convert Julia colon to Python slice(None)
function to_python_index(::Colon)
   return slice(nothing)
end

# Create figure and GridSpec layout
fig = plt.figure(figsize=(12, 10))
gs = gridspec.GridSpec(4, 2, figure=fig)

# Use Julia-style indexing with the wrapper
jgs = JulianGridSpec(gs)
ax1 = fig.add_subplot(jgs[1, 1:2])     # Row 1, columns 1-2
ax2 = fig.add_subplot(jgs[2:4, 1])     # Rows 2-4, column 1
module GridSpec

# Export the main struct and functions that users need
export JulianGridSpec, gridspec

using PyCall:pyimport, pycall, pybuiltin, PyObject

# Import matplotlib's gridspec module
const gridspec = pyimport("matplotlib.gridspec")

# Julia wrapper for Python's slice function
slice(args...) = pycall(pybuiltin("slice"), PyObject, args...)

# Wrapper struct for matplotlib GridSpec with Julia-style indexing
struct JulianGridSpec
    gs::PyObject  # Python GridSpec object
end

# Support for `end` keyword - implement lastindex
function Base.lastindex(jgs::JulianGridSpec, dim::Int)
    if dim == 1
        return jgs.gs.nrows
    elseif dim == 2
        return jgs.gs.ncols
    else
        error("GridSpec only has 2 dimensions")
    end
end

# Enable Julia-style indexing (1-based) on GridSpec objects
function Base.getindex(jgs::JulianGridSpec, i::Union{Int, UnitRange, Colon}, j::Union{Int, UnitRange, Colon})
    # Convert Julia indices to Python 0-based indices
    py_i = to_python_index(i)
    py_j = to_python_index(j)
    return get(jgs.gs, (py_i, py_j))
end

# Convert Julia integer to Python 0-based index
function to_python_index(idx::Int)
    return idx - 1
end

# Convert Julia range to Python slice
function to_python_index(idx::UnitRange)
    return slice(idx.start - 1, idx.stop)
end

# Convert Julia colon to Python slice(None)
function to_python_index(::Colon)
    return slice(nothing)
end

end # module