Manera de contornear el borde exterior de la región de cuadrícula seleccionada en Python

Question

Tengo el siguiente código:

import numpy as np
import matplotlib.pyplot as plt

x = np.linspace(-np.pi/2, np.pi/2, 30)
y = np.linspace(-np.pi/2, np.pi/2, 30)
x,y = np.meshgrid(x,y)

z = np.sin(x**2+y**2)[:-1,:-1]

fig,ax = plt.subplots()
ax.pcolormesh(x,y,z)

Lo que da esta imagen:

Ahora digamos que quiero resaltar el borde de ciertos cuadros de cuadrícula:

highlight = (z > 0.9)

Podría usar la función de contorno, pero esto daría como resultado un contorno "suavizado". Solo quiero resaltar el borde de una región, siguiendo el borde de los cuadros de la cuadrícula.

Lo más cerca que he llegado es añadiendo algo como esto:

highlight = np.ma.masked_less(highlight, 1)

ax.pcolormesh(x, y, highlight, facecolor = 'None', edgecolors = 'w')

Lo que da esta trama:

Lo cual está cerca, pero lo que realmente quiero es que solo se resalten los bordes exterior e interior de esa "rosquilla".

Básicamente, estoy buscando un híbrido de las funciones de contorno y pcolormesh, algo que sigue el contorno de algún valor, pero sigue los contenedores de cuadrícula en "pasos" en lugar de conectarse punto a punto. ¿Tiene sentido?

Nota al margen: en los argumentos de pcolormesh, tengo edgecolors = 'w', pero los bordes siguen siendo azules. ¿Que esta pasando ahí?

EDITAR: la respuesta inicial de JohanC usando add_iso_line() funciona para la pregunta tal como se planteó. Sin embargo, los datos reales que estoy usando son una cuadrícula x,y muy irregular, que no se puede convertir a 1D (como se requiere para add_iso_line().

Estoy usando datos que se han convertido de coordenadas polares (rho, phi) a cartesianas (x, y). La solución 2D planteada por JohanC no parece funcionar para el siguiente caso:

import numpy as np
import matplotlib.pyplot as plt
from scipy import ndimage

def pol2cart(rho, phi):
    x = rho * np.cos(phi)
    y = rho * np.sin(phi)
    return(x, y)

phi = np.linspace(0,2*np.pi,30)
rho = np.linspace(0,2,30)

pp, rr = np.meshgrid(phi,rho)

xx,yy = pol2cart(rr, pp)

z = np.sin(xx**2 + yy**2)

scale = 5
zz = ndimage.zoom(z, scale, order=0)

fig,ax = plt.subplots()
ax.pcolormesh(xx,yy,z[:-1, :-1])

xlim = ax.get_xlim()
ylim = ax.get_ylim()
xmin, xmax = xx.min(), xx.max()
ymin, ymax = yy.min(), yy.max()
ax.contour(np.linspace(xmin,xmax, zz.shape[1]) + (xmax-xmin)/z.shape[1]/2,
           np.linspace(ymin,ymax, zz.shape[0]) + (ymax-ymin)/z.shape[0]/2,
           np.where(zz < 0.9, 0, 1), levels=[0.5], colors='red')
ax.set_xlim(*xlim)
ax.set_ylim(*ylim)

Answer 1

Esta publicación muestra una forma de dibujar tales líneas. Como no es sencillo adaptarse al código actual pcolormesh, el siguiente código muestra una posible adaptación. Tenga en cuenta que se ha cambiado el nombre de las versiones 2d de xey, ya que se necesitan las versiones 1d para los segmentos de línea.

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection

x = np.linspace(-np.pi / 2, np.pi / 2, 30)
y = np.linspace(-np.pi / 2, np.pi / 2, 30)
xx, yy = np.meshgrid(x, y)

z = np.sin(xx ** 2 + yy ** 2)[:-1, :-1]

fig, ax = plt.subplots()
ax.pcolormesh(x, y, z)

def add_iso_line(ax, value, color):
    v = np.diff(z > value, axis=1)
    h = np.diff(z > value, axis=0)

    l = np.argwhere(v.T)
    vlines = np.array(list(zip(np.stack((x[l[:, 0] + 1], y[l[:, 1]])).T,
                               np.stack((x[l[:, 0] + 1], y[l[:, 1] + 1])).T)))
    l = np.argwhere(h.T)
    hlines = np.array(list(zip(np.stack((x[l[:, 0]], y[l[:, 1] + 1])).T,
                               np.stack((x[l[:, 0] + 1], y[l[:, 1] + 1])).T)))
    lines = np.vstack((vlines, hlines))
    ax.add_collection(LineCollection(lines, lw=1, colors=color))

add_iso_line(ax, 0.9, 'r')
plt.show()

Aquí hay una adaptación de la segunda respuesta, que puede funcionar solo con matrices 2d:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from scipy import ndimage

x = np.linspace(-np.pi / 2, np.pi / 2, 30)
y = np.linspace(-np.pi / 2, np.pi / 2, 30)
x, y = np.meshgrid(x, y)

z = np.sin(x ** 2 + y ** 2)

scale = 5
zz = ndimage.zoom(z, scale, order=0)

fig, ax = plt.subplots()
ax.pcolormesh(x, y,  z[:-1, :-1] )
xlim = ax.get_xlim()
ylim = ax.get_ylim()
xmin, xmax = x.min(), x.max()
ymin, ymax = y.min(), y.max()
ax.contour(np.linspace(xmin,xmax, zz.shape[1]) + (xmax-xmin)/z.shape[1]/2,
           np.linspace(ymin,ymax, zz.shape[0]) + (ymax-ymin)/z.shape[0]/2,
           np.where(zz < 0.9, 0, 1), levels=[0.5], colors='red')
ax.set_xlim(*xlim)
ax.set_ylim(*ylim)
plt.show()

Answer 2

Intentaré refactorizar el add_iso_linemétodo para que quede más claro y abierto a optimizaciones. Entonces, al principio, viene una parte obligada:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection

x = np.linspace(-np.pi/2, np.pi/2, 30)
y = np.linspace(-np.pi/2, np.pi/2, 30)
x, y = np.meshgrid(x,y)
z = np.sin(x**2+y**2)[:-1,:-1]

fig, ax = plt.subplots()
ax.pcolormesh(x,y,z)
xlim, ylim = ax.get_xlim(), ax.get_ylim()
highlight = (z > 0.9)

Ahora highlighthay una matriz binaria que se ve así:

después de eso, podemos extraer índices de celdas verdaderas, buscar vecindades falsas e identificar posiciones de líneas 'rojas'. No me siento lo suficientemente cómodo haciéndolo de una manera vectorizada (como aquí en el add_iso_linemétodo), así que solo uso un bucle simple:

lines = []
cells = zip(*np.where(highlight))
for x, y in cells:
    if x == 0 or highlight[x - 1, y] == 0: lines.append(([x, y], [x, y + 1]))
    if x == highlight.shape[0] or highlight[x + 1, y] == 0: lines.append(([x + 1, y], [x + 1, y + 1]))
    if y == 0 or highlight[x, y - 1] == 0: lines.append(([x, y], [x + 1, y]))
    if y == highlight.shape[1] or highlight[x, y + 1] == 0: lines.append(([x, y + 1], [x + 1, y + 1]))

Y, finalmente, cambio el tamaño y centro las coordenadas de las líneas para que encajen con pcolormesh:

lines = (np.array(lines) / highlight.shape - [0.5, 0.5]) * [xlim[1] - xlim[0], ylim[1] - ylim[0]]
ax.add_collection(LineCollection(lines, colors='r'))
plt.show()

En conclusión, esta es una solución muy similar a la de JohanC y, en general, más lenta. Afortunadamente, podemos reducir cellssignificativamente la cantidad extrayendo contornos solo usando el python-opencvpaquete:

import cv2
highlight = highlight.astype(np.uint8)
contours, hierarchy = cv2.findContours(highlight, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
cells = np.vstack(contours).squeeze()

Esta es una ilustración de las celdas que se están comprobando: