Removed plotly code; created function to check clustering in PyVISTA

random_seed: 333

check_clustering: True

G332.96: {

  n_clusters: 5,

  min_velocity: -68,

  surface_point: [291, 216, -66.43255],

  plot_velocity_limits: [-70, -62],

  plot_species: ['CH3OH'],

  continuum_limits: [0.1, 1]

}

G28: {

  n_clusters: 3,

  n_clusters: 2,

  min_velocity: -inf,

  surface_point: [272, 278, 79.2268463],

  plot_velocity_limits: [74, 81],

import pickle

import plotly.graph_objects as go

import plotly.express as px

import plotly.offline as poff

import matplotlib.pyplot as plt

import numpy as np

import matplotlib

import os

import yaml

import seaborn as sns

import pyspeckit

import astropy.units as u

matplotlib.use('TkAgg')

source_names = ['G28', ]

source_names = ['G332.96', ]

marker_opacity = 0.05

marker_size = 1.5

                                                       (velocity_array[:, 2] == config['surface_point'][1]) &

                                                       (np.isclose(velocity_array[:, 0],

                                                                   config['surface_point'][2]))])[0]

        # colors = ['blue', 'red', 'green', 'cyan', 'magenta']

        # fig = px.scatter_3d()

        for cluster in clusters:

            # get row indexes for samples with this cluster

            row_ix = tuple(np.nonzero(y_hat == cluster)[0])

            # fig.add_scatter3d(x=velocity_array[row_ix, 1],

            #                   y=velocity_array[row_ix, 2],

            #                   z=velocity_array[row_ix, 0],

            #                   mode='markers',

            #                   marker_color=colors[cluster],

            #                   marker_size=1)

            if cluster == surface_cluster:

                for val, iy, ix in velocity_array[row_ix, :]:

                    if val > config['min_velocity']:

    return surf, surface_cluster, y_hat

def check_clustering(velocity_array, n_clusters, cluster_assignments):

    ax = pv.Plotter()

    colors = ['blue', 'red', 'green', 'cyan', 'magenta']

    for cluster in range(n_clusters):

        row_ix = tuple(np.nonzero(cluster_assignments == cluster)[0])

        mesh = pv.PolyData(velocity_array[row_ix, :])

        mesh.points[:, -1] /= 10.0

        ax.add_mesh(mesh, render_points_as_spheres=True, opacity=0.2, color=colors[cluster])

        ax.show_grid(xtitle='dRA [pc]', ytitle='dDec [pc]', ztitle='Velocity', minor_ticks=True, show_zaxis=True)

    ax.show()

def plane_3d(xy: tuple,

             ra_grad: float,

             dec_grad: float,

                             (velocities <= config[source]['plot_velocity_limits'][1]))

            flattened_velocity_fields[species] = np.stack(

                [velocities, x_pix_flat, y_pix_flat, x_ax_flat, y_ax_flat, velocity_uncertainties], axis=1)

            # fig = px.scatter_3d()

            # fig.add_scatter3d(x=flattened_velocity_fields[:, 3][mask],

            #                   y=flattened_velocity_fields[:, 4][mask],

            #                   z=flattened_velocity_fields[:, 0][mask],

            #                   mode='markers')

            # fig.update_traces(marker=dict(size=marker_size, color=marker_color, opacity=marker_opacity, symbol='circle'))

            surf, surface_cluster, y_hat = get_surface_for_fitted_points(

                velocity_array=flattened_velocity_fields[species][:, 0:3][mask[species]],

                config=config[source],

                grid=grid,

                random_seed=config['random_seed']

            )

            if config['check_clustering'] is True:

                check_clustering(velocity_array=flattened_velocity_fields[species][:, (3, 4, 0)][mask[species]],

                                 n_clusters=config[source]['n_clusters'],

                                 cluster_assignments=y_hat)

            if ~np.isnan(surface_cluster):

                # fig.add_surface(x=x_ax, y=y_ax, z=surf)

                fitted_plane, fitted_plane_values, plane_pars_chain = fit_plane_to_gradient_with_uncertainties(

                    cluster_assignments=y_hat,

                    surface_cluster_id=surface_cluster,

                    velocity_uncertainties=flattened_velocity_fields[species][:, 5][mask[species]],

                    pixel_size_pc=pixel_size_pc

                )

                # fig.add_surface(x=x_ax, y=y_ax, z=fitted_plane_values,

                #                 opacity=0.2,

                #                 surfacecolor=['blue' for z in fitted_plane_values.flatten()])

                # The angle has been calculated using these points from the mom 0 map: p1: (145, 125) or (-0.48, -0.567) pc and p2: (330, 335) or (0.327, 0.349) pc; the centre was not changed, a possible improvement for the future

                rotated_plane = (

                    fitted_plane[0][0] * np.cos(source_properties[source]['rotation_angle'] / 360 * 2 * np.pi) +

                    species=species

                )

        # fig.show()

        make_3d_velocity_field_figure(flattened_velocity_fields, mask, source, x_ax, y_ax)

def make_3d_velocity_field_figure(flattened_velocity_fields, mask, source, x_ax, y_ax):

    mesh = {}

    ax = pv.Plotter()

    for species, color in zip(config[source]['plot_species'], ['cyan', 'red']):

        else:

            ax.add_mesh(mesh[species], render_points_as_spheres=True, opacity=0.05,

                        clim=config[source]['plot_velocity_limits'])

    hdu_cont = fits.open(os.path.join(datadir, continuum_fits[source]))[0]

    flattened_continuum_grid = np.stack((

        x_ax.flatten(),

    continuum_mesh = pv.PolyData(flattened_continuum_grid)

    continuum_mesh['Flux density [mJy/beam]'] = flattened_continuum_grid[:, 2] * 1000

    projected = continuum_mesh.project_points_to_plane(origin=(0, 0, ax.bounds[4]))

    ax.add_mesh(projected, cmap='hot', log_scale=True,

                clim=[config[source]['continuum_limits'][0], config[source]['continuum_limits'][1]])

    ax.show_grid(xtitle='dRA [pc]', ytitle='dDec [pc]', ztitle='Velocity', minor_ticks=True, show_zaxis=True,