NanoBrain-unitypackage/Editor/ClusterViewer.cs

617 lines
29 KiB
C#

using System.Collections.Generic;
using System.Linq;
using UnityEditor;
using UnityEngine;
using UnityEngine.UIElements;
namespace NanoBrain {
public class ClusterViewer : Editor {
public class GraphView : VisualElement {
protected readonly ClusterPrefab prefab;
protected SerializedObject serializedBrain;
protected Nucleus currentNucleus;
protected GameObject gameObject;
private List<NeuroidLayer> layers = new();
private readonly Dictionary<Nucleus, Vector2Int> neuroidPositions = new();
private bool expandArray = false;
protected ClusterPrefab prefabAsset;
protected VisualElement outputContainer;
protected readonly PopupField<string> outputsField;
public GraphView(ClusterPrefab prefab) {
this.prefab = prefab;
name = "content";
style.flexGrow = 1;
IMGUIContainer graphContainer = new(OnIMGUI);
graphContainer.style.position = Position.Absolute;
graphContainer.style.left = 0; graphContainer.style.top = 0;
graphContainer.style.right = 0; graphContainer.style.bottom = 0;
graphContainer.pickingMode = PickingMode.Position;
graphContainer.focusable = true;
Add(graphContainer);
outputContainer = new() {
style = {
flexDirection = FlexDirection.Row,
alignItems = Align.Center,
}
};
List<string> names = this.prefab.outputs.Select(output => output.name).ToList();
if (names.Count > 0 && names.First() != null) {
outputsField = new(names, names.First()) {
style = { flexGrow = 1 }
};
outputsField.RegisterValueChangedCallback(evt => OnOutputChanged(evt.newValue));
outputContainer.Add(outputsField);
}
Add(outputContainer);
// Subscribe when added to panel (editor UI ready)
RegisterCallback<AttachToPanelEvent>(evt => Subscribe());
RegisterCallback<DetachFromPanelEvent>(evt => Unsubscribe());
}
void OnOutputChanged(string outputName) {
if (this.currentNucleus.parent != null)
// Get nucleus in the parent instance
this.currentNucleus = this.currentNucleus.parent.GetNucleus(outputName);
else
// Get nucleus in the prefab
this.currentNucleus = this.prefab.GetNucleus(outputName);
}
bool subscribed = false;
void Subscribe() {
if (subscribed) return;
SceneView.duringSceneGui += OnSceneGUI;
subscribed = true;
SceneView.RepaintAll();
}
void Unsubscribe() {
if (!subscribed) return;
SceneView.duringSceneGui -= OnSceneGUI;
subscribed = false;
}
public void SetGraph(GameObject gameObject, Nucleus nucleus) { //}, VisualElement inspectorContainer) {
this.gameObject = gameObject;
//this.cluster = brain;
if (Application.isPlaying == false)
this.serializedBrain = new SerializedObject(this.prefab);
this.currentNucleus = nucleus;
Rebuild(); //inspectorContainer);
}
void Rebuild() { //VisualElement inspectorContainer) {
BuildLayers();
if (this.currentNucleus == null) {
// inspectorContainer.Clear();
return;
}
string path = AssetDatabase.GetAssetPath(this.prefab); // or known path
this.prefabAsset = AssetDatabase.LoadAssetAtPath<ClusterPrefab>(path);
if (this.prefabAsset == null) {
// create in memory save if it doesn't exist
this.prefabAsset = CreateInstance<ClusterPrefab>();
//Debug.LogError("Cluster Prefab is not found on disk");
}
//DrawInspector(inspectorContainer);
}
protected void BuildLayers() {
// A temporary list to track what's been added to layers
this.layers = new();
int layerIx = 0;
Nucleus selectedNucleus = this.currentNucleus;
if (selectedNucleus == null)
return;
NeuroidLayer currentLayer = new() { ix = layerIx };
if (selectedNucleus is Neuron selectedNeuron && selectedNeuron.receivers != null) {
foreach (Nucleus receiver in selectedNeuron.receivers) {
Nucleus outputNeuroid = receiver;
if (outputNeuroid != null) {
AddToLayer(currentLayer, outputNeuroid);
// Debug.Log($"layer {layerIx} nucleus {outputNeuroid.name}");
}
}
}
if (currentLayer.neuroids.Count > 0) {
this.layers.Add(currentLayer);
layerIx++;
currentLayer = new() { ix = layerIx };
}
AddToLayer(currentLayer, selectedNucleus);
this.layers.Add(currentLayer);
// Debug.Log($"layer {layerIx} nucleus {selectedNucleus.name}");
layerIx++;
currentLayer = new() { ix = layerIx };
if (selectedNucleus.synapses != null) {
foreach (Synapse synapse in selectedNucleus.synapses) {
Nucleus input = synapse.neuron;
AddToLayer(currentLayer, input);
// Debug.Log($"layer {layerIx} nucleus {input.name}");
}
}
if (currentLayer.neuroids.Count > 0) {
this.layers.Add(currentLayer);
}
}
private void AddToLayer(NeuroidLayer layer, Nucleus nucleus) {
if (nucleus == null)
return;
layer.neuroids.Add(nucleus);
//nucleus.layerIx = layer.ix;
// Store its position
Vector2Int neuroidPosition = new(layer.ix, layer.neuroids.Count - 1);
neuroidPositions[nucleus] = neuroidPosition;
}
public void OnIMGUI() {
if (currentNucleus == null)
return;
if (Application.isPlaying == false)
serializedBrain.Update();
Handles.BeginGUI();
DrawGraph();
Handles.EndGUI();
}
private void DrawGraph() {
float size = 20;
Vector3 position = new(150, 210, 0);
DrawReceivers(this.currentNucleus, position, size);
DrawSynapses(this.currentNucleus, position, size);
// Draw selected Nucleus
if (expandArray) {
// if (this.currentNucleus is IReceptor receptor1) {
// float maxValue = 0;
// foreach (Nucleus nucleus in receptor1.nucleiArray) {
// if (nucleus is Neuron neuron) {
// float value = neuron.outputMagnitude;
// if (value > maxValue)
// maxValue = value;
// }
// }
// float spacing = 400f / receptor1.nucleiArray.Count();
// float margin = 10 + spacing / 2;
// float xMin = 150 - size;
// float xMax = 150 + size;
// float yMin = 10 + margin - size / 2;
// float yMax = 400 - margin + size;
// Vector3[] verts = new Vector3[4] {
// new(xMin, yMin, 0),
// new(xMax, yMin, 0),
// new(xMax, yMax, 0),
// new(xMin, yMax, 0)
// };
// Handles.color = Color.black;
// Handles.DrawAAConvexPolygon(verts);
// int row = 0;
// foreach (Nucleus nucleus in receptor1.nucleiArray) {
// Vector3 pos = new(150, margin + row * spacing, 0.0f);
// Handles.color = Color.white;
// // The selected nucleus highlight ring
// Handles.DrawSolidDisc(pos, Vector3.forward, size + 2);
// DrawNucleus(nucleus, pos, maxValue, size);
// row++;
// }
// GUIStyle style = new(EditorStyles.label) {
// alignment = TextAnchor.UpperCenter,
// normal = { textColor = Color.white },
// fontStyle = FontStyle.Bold,
// };
// Vector3 labelPos = new(150, yMax + size + 5, 0);
// string receptorName = receptor1.GetName();
// int colonPos = receptorName.IndexOf(":");
// if (colonPos > 0) {
// string baseName = receptorName[..colonPos];
// Handles.Label(labelPos, baseName, style);
// }
// else
// Handles.Label(labelPos, receptorName, style);
// }
// else {
Handles.color = Color.white;
// The selected nucleus highlight ring
Handles.DrawSolidDisc(position, Vector3.forward, size + 2);
float maxValue = 1;
if (this.currentNucleus is Neuron neuron)
maxValue = neuron.outputMagnitude;
else if (this.currentNucleus is Cluster cluster)
maxValue = cluster.defaultOutput.outputMagnitude;
DrawNucleus(this.currentNucleus, position, maxValue, 20);
// }
}
else {
Handles.color = Color.white;
// The selected nucleus highlight ring
Handles.DrawSolidDisc(position, Vector3.forward, size + 2);
float maxValue = 1;
if (this.currentNucleus is Neuron neuron)
maxValue = neuron.outputMagnitude;
else if (this.currentNucleus is Cluster cluster)
maxValue = cluster.defaultOutput.outputMagnitude;
DrawNucleus(this.currentNucleus, position, maxValue, 20);
}
}
private void DrawReceivers(Nucleus nucleus, Vector3 parentPos, float size) {
List<Nucleus> receivers;
if (nucleus is Neuron neuron)
receivers = neuron.receivers;
else if (nucleus is Cluster cluster)
receivers = cluster.CollectReceivers();
else
return;
int nodeCount = receivers.Count(); //neuron != null ? neuron.receivers.Count() : 1;
// Determine the maximum value in this layer
// This is used to 'scale' the output value colors of the nuclei
float maxValue = 0;
foreach (Nucleus receiver in receivers) {
if (receiver is Neuron neuroid) {
float value = neuroid.outputMagnitude;
if (value > maxValue)
maxValue = value;
}
}
// Determine the spacing of the nuclei in the layer
float spacing = 400f / nodeCount;
float margin = 10 + spacing / 2;
int row = 0;
List<Nucleus[]> drawnArrays = new();
foreach (Nucleus receiver in receivers) {
// if (receiver is Receptor receptor) {
// if (drawnArrays.Contains(receptor.nucleiArray))
// continue;
// drawnArrays.Add(receptor.nucleiArray);
// }
Nucleus receiverNucleus = receiver;
if (receiverNucleus == null)
continue;
Vector3 pos = new(50, margin + row * spacing, 0.0f);
Handles.color = Color.white;
Handles.DrawLine(parentPos, pos);
DrawNucleus(receiverNucleus, pos, maxValue, size);
row++;
}
}
private void DrawSynapses(Nucleus nucleus, Vector3 parentPos, float size) {
int nodeCount = nucleus.synapses.Count;
// Determine the maximum value in this layer
// This is used to 'scale' the output value colors of the nuclei
float maxValue = 0;
int neuronCount = 0;
List<Nucleus[]> drawnArrays = new();
foreach (Synapse synapse in nucleus.synapses) {
if (synapse.neuron == null)
continue;
// if (synapse.neuron is Receptor receptor) {
// if (drawnArrays.Contains(receptor.nucleiArray))
// continue;
// drawnArrays.Add(receptor.nucleiArray);
// }
// else if (synapse.neuron.parent is ClusterReceptor clusterReceptor) {
// if (drawnArrays.Contains(clusterReceptor.nucleiArray))
// continue;
// drawnArrays.Add(clusterReceptor.nucleiArray);
// }
if (synapse.neuron.parent is Cluster cluster && cluster.siblingClusters != null) {
if (drawnArrays.Contains(cluster.siblingClusters))
continue;
drawnArrays.Add(cluster.siblingClusters);
}
if (synapse.neuron is Neuron synapseNeuron) {
float value = synapseNeuron.outputMagnitude * synapse.weight;
// Debug.Log($"{synapse.nucleus.name}: {value} {length(synapse.nucleus.outputValue)} {synapse.weight}");
if (value > maxValue)
maxValue = value;
}
neuronCount++;
}
// Determine the spacing of the nuclei in the layer
float spacing = 400f / neuronCount;
float margin = 10 + spacing / 2;
int row = 0;
drawnArrays = new();
foreach (Synapse synapse in nucleus.synapses) {
if (synapse.neuron is null)
continue;
// if (synapse.neuron is Receptor neuron) {
// if (drawnArrays.Contains(neuron.nucleiArray))
// continue;
// drawnArrays.Add(neuron.nucleiArray);
// }
// else if (synapse.neuron.parent is ClusterReceptor clusterReceptor) {
// if (drawnArrays.Contains(clusterReceptor.nucleiArray))
// continue;
// drawnArrays.Add(clusterReceptor.nucleiArray);
// }
Vector3 pos = new(250, margin + row * spacing, 0.0f);
Handles.color = Color.white;
Handles.DrawLine(parentPos, pos);
Color color = Color.black;
if (Application.isPlaying) {
if (maxValue == 0 || !float.IsFinite(maxValue))
maxValue = 1;
float brightness = 0;
if (synapse.neuron is Neuron synapseNeuron)
brightness = synapseNeuron.outputMagnitude * synapse.weight / maxValue;
color = new Color(brightness, brightness, brightness, 1f);
}
if (synapse.neuron.parent != null && synapse.neuron.parent != this.currentNucleus.parent) {
// the synapse nucleus is part of a subcluster
//DrawNucleus(synapse.neuron.parent, pos, maxValue, size, color);
DrawNucleus(synapse.neuron, pos, maxValue, size, color);
}
else {
DrawNucleus(synapse.neuron, pos, maxValue, size, color);
}
row++;
}
}
private void DrawNucleus(Nucleus nucleus, Vector3 position, float maxValue, float size) {
Color color;
if (Application.isPlaying) {
float brightness = 0;
if (nucleus is Neuron neuron)
brightness = neuron.outputMagnitude / maxValue;
color = new Color(brightness, brightness, brightness, 1f);
}
else
color = Color.black;
DrawNucleus(nucleus, position, maxValue, size, color);
}
private void DrawNucleus(Nucleus nucleus, Vector3 position, float maxValue, float size, Color color) {
if (nucleus is MemoryCell) {
Handles.color = Color.white;
Handles.DrawWireDisc(position + Vector3.right * 10, Vector3.forward, size);
}
Handles.color = color;
Handles.DrawSolidDisc(position, Vector3.forward, size);
Handles.color = Color.white;
// Position the label in front of the disc
Vector3 labelPosition = position + (Vector3.forward * 0.1f);
GUIStyle style = new(EditorStyles.label) {
alignment = TextAnchor.MiddleCenter,
normal = { textColor = Color.white },
fontStyle = FontStyle.Bold,
};
// if (nucleus is IReceptor receptor1) {
// if (expandArray) {
// // Put array indices above elements
// style.alignment = TextAnchor.LowerCenter;
// Vector3 labelPos1 = position + Vector3.down * (size + 5); // below disc
// int colonPos1 = nucleus.name.IndexOf(":");
// if (colonPos1 > 0) {
// string extName = nucleus.name[(colonPos1 + 2)..];
// Handles.Label(labelPos1, extName, style);
// }
// }
// else {
// // draw the array size label
// if (color.grayscale > 0.5f)
// style.normal.textColor = Color.black;
// else
// style.normal.textColor = Color.white;
// Handles.Label(labelPosition, receptor1.nucleiArray.Length.ToString(), style);
// style.normal.textColor = Color.white;
// }
// }
// else
if (nucleus.parent != null && nucleus.parent is Cluster parentCluster) {
if (expandArray) {
// Put array indices above elements
style.alignment = TextAnchor.LowerCenter;
Vector3 labelPos1 = position + Vector3.down * (size + 5); // below disc
int colonPos1 = nucleus.name.IndexOf(":");
if (colonPos1 > 0) {
string extName = nucleus.name[(colonPos1 + 2)..];
Handles.Label(labelPos1, extName, style);
}
}
else {
if (parentCluster.siblingClusters != null && parentCluster.siblingClusters.Length > 1) {
// draw the array size label
if (color.grayscale > 0.5f)
style.normal.textColor = Color.black;
else
style.normal.textColor = Color.white;
Handles.Label(labelPosition, parentCluster.siblingClusters.Length.ToString(), style);
style.normal.textColor = Color.white;
}
}
}
else if (nucleus is Cluster cluster) {
if (expandArray) {
// Put array indices above elements
style.alignment = TextAnchor.LowerCenter;
Vector3 labelPos1 = position + Vector3.down * (size + 5); // below disc
int colonPos1 = nucleus.name.IndexOf(":");
if (colonPos1 > 0) {
string extName = nucleus.name[(colonPos1 + 2)..];
Handles.Label(labelPos1, extName, style);
}
}
else {
if (cluster.siblingClusters != null && cluster.siblingClusters.Length > 1) {
// draw the array size label
if (color.grayscale > 0.5f)
style.normal.textColor = Color.black;
else
style.normal.textColor = Color.white;
Handles.Label(labelPosition, cluster.siblingClusters.Length.ToString(), style);
style.normal.textColor = Color.white;
}
}
}
if (expandArray == false) {// || nucleus is not IReceptor) {
// put name below nucleus
Vector3 labelPos = position - Vector3.down * (size + 5); // below neuron
style.alignment = TextAnchor.UpperCenter;
if (nucleus.parent != null && nucleus.parent is Cluster parentCluster1) {
parentCluster1.name ??= "";
string baseName = "";
if (parentCluster1 != currentNucleus.parent) {
int colonPos = parentCluster1.name.IndexOf(":");
if (colonPos > 0 && colonPos < parentCluster1.name.Length - 2)
baseName = parentCluster1.name[..colonPos] + ".";
else
baseName = parentCluster1.name + ".";
}
// if (colonPos > 0 && colonPos < parentCluster1.name.Length - 2) {
// // if it is an array, we should not show the :0 of the first element
// //baseName = baseName[..colonPos];
// Handles.Label(labelPos, baseName + nucleus.name, style);
// }
// else
Handles.Label(labelPos, baseName + nucleus.name, style);
}
else {
nucleus.name ??= "";
int colonPos = nucleus.name.IndexOf(":");
if (colonPos > 0 && colonPos < nucleus.name.Length - 2) {
// if it is an array, we should not show the :0 of the first element
string baseName = nucleus.name[..colonPos];
Handles.Label(labelPos, baseName, style);
}
else
Handles.Label(labelPos, nucleus.name, style);
}
}
// Draw Cluster ring
if (nucleus.parent != currentNucleus.parent || nucleus is Cluster) {
Handles.color = Color.white;
Handles.DrawWireDisc(position, Vector3.forward, size + 5);
}
// Tooltip
Rect neuronRect = new(position.x - size, position.y - size, size * 2, size * 2);
int id = GUIUtility.GetControlID(FocusType.Passive);
Event e = Event.current;
EventType et = e.GetTypeForControl(id);
if (e != null && neuronRect.Contains(e.mousePosition)) {
// Process Hover
HandleMouseHover(nucleus, neuronRect);
// Process click
if (e.type == EventType.MouseDown && e.button == 0) {
// Consume the event so the scene doesn't also handle it
e.Use();
if (nucleus.parent != null && nucleus.parent is Cluster parentCluster2)
HandleClicked(parentCluster2);
else
HandleClicked(nucleus);
}
}
}
private void HandleMouseHover(Nucleus nucleus, Rect rect) {
GUIContent tooltip;
if (nucleus is Neuron neuron) {
tooltip = new(
$"{nucleus.name}" +
$"\nValue: {neuron.outputMagnitude}");
}
else
tooltip = new($"{nucleus.name}");
Vector2 mousePosition = Event.current.mousePosition;
// Display tooltip with some offset
Vector2 tooltipSize = GUI.skin.box.CalcSize(tooltip);
Rect tooltipRect = new Rect(mousePosition.x + 10, mousePosition.y + 10, tooltipSize.x, tooltipSize.y);
GUI.Box(tooltipRect, tooltip);
}
private void HandleClicked(Nucleus nucleus) {
if (nucleus == this.currentNucleus) {
if (nucleus is Cluster) //is Receptor) // || nucleus is ClusterReceptor)
expandArray = !expandArray;
else
expandArray = false;
}
else {
this.currentNucleus = nucleus;
expandArray = false;
BuildLayers();
}
}
void OnSceneGUI(SceneView sceneView) {
if (this.gameObject != null) {
// if (this.currentNucleus is IReceptor receptor) {
// foreach (Nucleus nucleus in receptor.nucleiArray) {
// if (nucleus is Neuron neuron) {
// Vector3 worldVector = this.gameObject.transform.TransformVector(neuron.outputValue);
// Handles.color = Color.yellow;
// Handles.DrawLine(this.gameObject.transform.position, this.gameObject.transform.position + worldVector);
// }
// }
// }
// else {
if (this.currentNucleus is Neuron currentNeuron) {
Vector3 worldVector = this.gameObject.transform.TransformVector(currentNeuron.outputValue);
Handles.color = Color.yellow;
Handles.DrawLine(this.gameObject.transform.position, this.gameObject.transform.position + worldVector);
}
// }
}
}
}
}
public class NeuroidLayer {
public int ix = 0;
public List<Nucleus> neuroids = new();
}
}