690 lines
29 KiB
C#

using System.Collections.Generic;
using UnityEngine;
using UnityEditor;
namespace NanoBrain.Unity {
public class ClusterView {
private static readonly float discRadius = 20;
static readonly Dictionary<string, ClusterView> viewStates = new();
public static ClusterView GetClusterView(SerializedProperty property) {
string key = property.propertyPath + "_" + property.serializedObject.targetObject.GetEntityId();
if (!viewStates.TryGetValue(key, out ClusterView state))
state = new() { key = key };
return state;
}
public static ClusterView GetClusterView(SerializedObject serializedObject) {
string key = serializedObject.targetObject.GetEntityId().ToString();
if (!viewStates.TryGetValue(key, out ClusterView state))
state = new() { key = key };
return state;
}
private void UpdateViewState() {
viewStates[this.key] = this;
}
public static void Render(Rect drawRect, Cluster cluster, SerializedProperty property) {
ClusterView clusterView = GetClusterView(property);
clusterView.currentCluster ??= cluster;
clusterView.Render(drawRect);
}
public static void Render(Rect drawRect, Cluster cluster, SerializedObject obj) {
ClusterView clusterView = GetClusterView(obj);
clusterView.currentCluster ??= cluster;
clusterView.Render(drawRect);
}
public void Render(Rect drawRect) {
// background
EditorGUI.DrawRect(drawRect, Color.black);
const float contentWidth = 1000f;
Rect contentRect = new(0f, 0f, contentWidth, drawRect.height - 20);
// Begin horizontal-only scroll view
this.scrollPos = GUI.BeginScrollView(drawRect, this.scrollPos, contentRect, false, false);
// Local content group: draw GUI content using content-local coords (0..contentWidth)
GUI.BeginGroup(new Rect(-this.scrollPos.x, 0f, contentWidth, drawRect.height));
EditorGUI.DrawRect(new Rect(0f, 0f, contentWidth, drawRect.height), new Color(0.08f, 0.08f, 0.08f, 1f));
GUI.EndGroup();
GUI.EndScrollView();
// Clip to drawRect so Handles are not drawn outside the black area
GUI.BeginGroup(drawRect);
// Inner group positions content origin so local coords match content space and respect scroll
GUI.BeginGroup(new Rect(-this.scrollPos.x, 0f, contentWidth, drawRect.height));
Handles.BeginGUI();
this.DrawFocusGraph();
Handles.EndGUI();
GUI.EndGroup(); // end inner group
GUI.EndGroup(); // end clipping group
UpdateViewState();
}
public string key = null;
public Vector2 scrollPos = Vector2.zero;
public bool expandArray = false;
public Cluster currentCluster;
public Nucleus currentNucleus = null;
public Nucleus selectedSynapseNeuron = null;
public Nucleus selectedOutput;
protected void DrawFocusGraph() {
float size = 20;
Vector3 position = new(150, 210, 0);
if (this.currentNucleus != null) {
DrawReceivers(this.currentNucleus, position, size);
DrawSynapses(this.currentNucleus, position, size);
// Draw selected Nucleus
if (this.expandArray) {
float maxValue = 1;
if (this.currentNucleus is Cluster cluster) {
float spacing = 400f / cluster.instanceCount;
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;
if (cluster.instances == null) {
Vector2 pos = new(150, margin + row * spacing);
Handles.color = Color.white;
// The selected sibling highlight ring
Handles.DrawSolidDisc(pos, Vector3.forward, size + 2);
DrawNucleus(cluster, pos, maxValue);
row++;
}
else {
foreach (Cluster sibling in cluster.instances) {
Vector3 pos = new(150, margin + row * spacing, 0.0f);
Handles.color = Color.white;
// The selected sibling highlight ring
Handles.DrawSolidDisc(pos, Vector3.forward, size + 2);
DrawNucleus(sibling, pos, maxValue);
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 clusterName = cluster.name;
int colonPos = clusterName.IndexOf(":");
if (colonPos > 0) {
string baseName = clusterName[..colonPos];
Handles.Label(labelPos, baseName, style);
}
else
Handles.Label(labelPos, clusterName, style);
}
else {
if (this.currentNucleus is Neuron neuron)
maxValue = neuron.outputMagnitude;
DrawNucleus(this.currentNucleus, position, maxValue);
}
}
else {
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);
}
}
else {
DrawAllOutputs(position);
DrawOutputs(position);
}
}
protected 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(true);
else
return;
// For top-level nodes, add link to previous editor and/or 'Outputs'
int nodeCount = receivers.Count;
if (nucleus == this.selectedOutput) {
// Add link to 'Outpus'
nodeCount++;
if (ClusterViewer.previousPrefab != null)
// Add link to previous editor
nodeCount++;
}
// 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;
foreach (Nucleus receiver in receivers) {
Nucleus receiverNucleus = receiver;
if (receiverNucleus == null)
continue;
Vector3 pos = new(50, margin + row * spacing, 0.0f);
DrawEdge(parentPos, pos);
DrawNucleus(receiverNucleus, pos, maxValue);
row++;
}
if (nucleus == this.selectedOutput) {
Vector3 pos = new(50, margin + row * spacing, 0);
if (ClusterViewer.previousPrefab != null) {
DrawEdge(parentPos, pos);
DrawClusterPrefab(ClusterViewer.previousPrefab, pos);
row++;
}
pos = new(50, margin + row * spacing, 0);
DrawEdge(parentPos, pos);
DrawAllOutputs(pos);
}
}
protected void DrawSynapses(Nucleus nucleus, Vector3 parentPos, float size) {
if (nucleus is not Neuron neuron)
return;
if (this.selectedSynapseNeuron != null) {
DrawClusterSynapses(this.selectedSynapseNeuron, parentPos);
return;
}
if (nucleus == null)
return;
// 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<string> drawnNeuronNames = new();
foreach (Synapse synapse in neuron.synapses) {
if (synapse.neuron == null)
continue;
// Count multiple synapses to the same neuron only once
string neuronName = synapse.neuron.name;
if (synapse.neuron.parent != null)
neuronName = synapse.neuron.parent.baseName + "." + neuronName;
if (drawnNeuronNames.Contains(neuronName))
continue;
drawnNeuronNames.Add(neuronName);
float value = synapse.neuron.outputMagnitude * 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;
//List<Neuron> drawnNeurons = new();
drawnNeuronNames = new();
foreach (Synapse synapse in neuron.synapses) {
if (synapse.neuron is null)
continue;
// Draw multiple synapses to the same neuron only once
string neuronName = synapse.neuron.name;
if (synapse.neuron.parent != null)
neuronName = synapse.neuron.parent.baseName + "." + neuronName;
if (drawnNeuronNames.Contains(neuronName))
continue;
drawnNeuronNames.Add(neuronName);
Vector3 pos = new(250, margin + row * spacing, 0.0f);
DrawEdge(parentPos, pos);
// 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 = synapse.neuron.outputMagnitude * synapse.weight / maxValue;
color = new Color(brightness, brightness, brightness, 1f);
}
DrawNucleus(synapse.neuron, pos, color);
row++;
}
}
protected void DrawClusterSynapses(Nucleus nucleus, Vector3 parentPos) {
if (nucleus == null || nucleus.parent == null || nucleus.parent.instances == null)
return;
// Hack to disable showing labels
expandArray = true;
float maxValue = 0;
foreach (Cluster sibling in nucleus.parent.instances) {
Neuron siblingNeuron = sibling.GetNucleus(nucleus.name) as Neuron;
float value = siblingNeuron.outputMagnitude; // no need to add weight as they are all the same
if (value > maxValue)
maxValue = value;
}
// Determine the spacing of the nuclei in the layer
float spacing = 400f / nucleus.parent.instanceCount;
float margin = 10 + spacing / 2;
int row = 0;
foreach (Cluster sibling in nucleus.parent.instances) {
Neuron siblingNeuron = sibling.GetNucleus(nucleus.name) as Neuron;
Vector3 position = new(250, margin + row * spacing, 0.0f);
DrawEdge(parentPos, position);
Color color = Color.black;
if (Application.isPlaying) {
if (maxValue == 0 || !float.IsFinite(maxValue))
maxValue = 1;
float brightness = siblingNeuron.outputMagnitude / maxValue;
color = new Color(brightness, brightness, brightness, 1f);
}
DrawNucleus(siblingNeuron, position, color);
GUIStyle style = new(EditorStyles.label) {
alignment = TextAnchor.UpperCenter,
normal = { textColor = Color.white },
fontStyle = FontStyle.Bold,
};
Vector3 labelPos = position - Vector3.down * (discRadius + 5); // below neuron
string name = $"{sibling.baseName}\n{nucleus.name}";
Handles.Label(labelPos, name, style);
row++;
}
expandArray = false;
}
protected void DrawOutputs(Vector2 parentPos) {
if (this.currentCluster == null)
return;
// 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> drawnNuclei = new();
foreach (Nucleus nucleus in this.currentCluster.outputs) {
if (nucleus is not Neuron neuron)
continue;
// Draw multiple synapses to the same neuron only once
if (drawnNuclei.Contains(nucleus))
continue;
drawnNuclei.Add(nucleus);
float value = neuron.outputMagnitude;
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;
drawnNuclei = new();
foreach (Nucleus nucleus in this.currentCluster.outputs) {
if (nucleus is not Neuron neuron)
continue;
// Draw multiple synapses to the same neuron only once
if (drawnNuclei.Contains(nucleus))
continue;
drawnNuclei.Add(nucleus);
Vector3 pos = new(250, margin + row * spacing, 0.0f);
DrawEdge(parentPos, pos);
Color color = Color.black;
if (Application.isPlaying) {
if (maxValue == 0 || !float.IsFinite(maxValue))
maxValue = 1;
float brightness = neuron.outputMagnitude / maxValue;
color = new Color(brightness, brightness, brightness, 1f);
}
DrawNucleus(nucleus, pos, color);
row++;
}
}
protected void DrawNucleus(Nucleus nucleus, Vector3 position, float maxValue) {
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, color);
}
protected void DrawNucleus(Nucleus nucleus, Vector2 position, Color color) {
if (nucleus == null)
return;
if (nucleus == this.currentNucleus) {
// The selected nucleus highlight ring
Handles.color = Color.white;
Handles.DrawSolidDisc(position, Vector3.forward, discRadius + 2);
}
if (nucleus is MemoryCell) {
Handles.color = Color.white;
Handles.DrawWireDisc(position + Vector2.right * 10, Vector3.forward, discRadius);
}
Handles.color = color;
Handles.DrawSolidDisc(position, Vector3.forward, discRadius);
Handles.color = Color.white;
// Position the label in front of the disc
//Vector3 labelPosition = position; // + (Vector2.forward * 0.1f);
GUIStyle style = new(EditorStyles.label) {
alignment = TextAnchor.MiddleCenter,
normal = { textColor = Color.white },
fontStyle = FontStyle.Bold,
};
if (nucleus.parent is Cluster parentCluster && this.currentNucleus != null && parentCluster != this.currentNucleus.parent)
DrawCluster(parentCluster, position, color);
else if (nucleus is Cluster cluster)
DrawCluster(cluster, position, color);
if (this.expandArray == false) {// || nucleus != currentNucleus) {
// put name below nucleus
Vector3 labelPos = position - Vector2.down * (discRadius + 5); // below neuron
style.alignment = TextAnchor.UpperCenter;
if (nucleus.parent != null && this.currentNucleus != null && nucleus.parent != this.currentNucleus.parent && nucleus.parent is Cluster parentCluster1) {
// This neuron is part of another cluster
parentCluster1.name ??= "";
int colonPos = parentCluster1.name.IndexOf(":");
string baseName;
if (colonPos > 0 && colonPos < parentCluster1.name.Length - 2)
baseName = parentCluster1.name[..colonPos] + "\n";
else
baseName = parentCluster1.name + "\n";
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);
}
}
// Tooltip
Rect neuronRect = new(position.x - discRadius, position.y - discRadius, discRadius * 2, discRadius * 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 is Cluster parentCluster2)
OnNeuronClick(parentCluster2);
else
OnNeuronClick(nucleus);
}
}
}
protected void DrawCluster(Cluster cluster, Vector3 position, Color color) {
GUIStyle labelTextStyle = new(EditorStyles.label) {
normal = { textColor = Color.white },
fontStyle = FontStyle.Bold,
};
if (this.expandArray) {
// Put array indices above the discs
labelTextStyle.alignment = TextAnchor.LowerCenter;
Vector3 labelPosition = position + Vector3.down * (discRadius + 5); // below disc
// Strip the instance number in the name
int colonPos1 = cluster.name.IndexOf(":");
if (colonPos1 > 0) {
string extName = cluster.name[(colonPos1 + 2)..];
Handles.Label(labelPosition, extName, labelTextStyle);
}
else
Handles.Label(labelPosition, "0", labelTextStyle);
}
else {
// Put instance count inside the disc
labelTextStyle.alignment = TextAnchor.MiddleCenter;
Vector3 labelPosition = position + (Vector3.forward * 0.1f);
// Adjust text color based on disc color
if (color.grayscale > 0.5f)
labelTextStyle.normal.textColor = Color.black;
else
labelTextStyle.normal.textColor = Color.white;
if (cluster.instanceCount > 1) {
Handles.Label(labelPosition, cluster.instanceCount.ToString(), labelTextStyle);
labelTextStyle.normal.textColor = Color.white;
}
else if (cluster.instances != null && cluster.instances.Length > 1) {
Handles.Label(labelPosition, cluster.instances.Length.ToString(), labelTextStyle);
labelTextStyle.normal.textColor = Color.white;
}
}
// Draw a circle around the disc to indicate this is a Cluster
Handles.color = Color.white;
Handles.DrawWireDisc(position, Vector3.forward, discRadius + 5);
}
protected void DrawClusterPrefab(ClusterPrefab prefab, Vector2 position) {
Handles.color = Color.black;
Handles.DrawSolidDisc(position, Vector3.forward, discRadius);
// Draw a circle around the disc to indicate this is a Cluster
Handles.color = Color.white;
Handles.DrawWireDisc(position, Vector3.forward, discRadius + 5);
// put name below nucleus
GUIStyle style = new(EditorStyles.label) {
alignment = TextAnchor.MiddleCenter,
normal = { textColor = Color.white },
fontStyle = FontStyle.Bold,
};
Vector2 labelPos = position - Vector2.down * (discRadius + 5); // below neuron
style.alignment = TextAnchor.UpperCenter;
Handles.Label(labelPos, prefab.name, style);
Rect neuronRect = new(position.x - discRadius, position.y - discRadius, discRadius * 2, discRadius * 2);
int id = GUIUtility.GetControlID(FocusType.Passive);
Event e = Event.current;
EventType et = e.GetTypeForControl(id);
if (e != null && neuronRect.Contains(e.mousePosition)) {
// Process click
if (e.type == EventType.MouseDown && e.button == 0) {
// Consume the event so the scene doesn't also handle it
e.Use();
Selection.activeObject = prefab;
EditorGUIUtility.PingObject(prefab);
ClusterViewer.previousPrefab = null;
Editor.CreateEditor(prefab);
}
}
}
protected void DrawAllOutputs(Vector2 position) {
GUIStyle labelTextStyle = new(EditorStyles.label) {
normal = { textColor = Color.white },
fontStyle = FontStyle.Bold,
alignment = TextAnchor.MiddleCenter,
};
Handles.Label(position, "Outputs", labelTextStyle);
Rect neuronRect = new(position.x - discRadius, position.y - discRadius, discRadius * 2, discRadius * 2);
Event e = Event.current;
if (e != null && neuronRect.Contains(e.mousePosition)) {
// Process click
if (e.type == EventType.MouseDown && e.button == 0) {
// Consume the event so the scene doesn't also handle it
e.Use();
OnAllOutputsClick();
}
}
}
protected void DrawEdge(Vector2 from, Vector2 to) {
Handles.color = Color.white;
// Handles.DrawLine(from, to);
Vector2 dir = to - from;
float len = dir.magnitude;
if (len <= 2f * discRadius || len <= Mathf.Epsilon)
// line too short
return;
Vector2 n = dir / len; // normalized
Vector2 a = from + n * discRadius;
Vector2 b = to - n * discRadius;
Handles.DrawLine(a, b);
}
#region Interaction
protected static 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);
}
protected void OnNeuronClick(Nucleus nucleus) {
if (nucleus == this.currentNucleus) {
this.selectedSynapseNeuron = null;
// if (Application.isPlaying) {
// if (nucleus is Cluster)
// expandArray = !expandArray;
// else
// expandArray = false;
// }
// else {
if (nucleus is Cluster cluster)
OnClusterClick(cluster);
// }
}
else if (nucleus.parent != null && this.currentNucleus != null && nucleus.parent != this.currentNucleus.parent) {
// We go to a different cluster
if (Application.isPlaying) {
if (this.selectedSynapseNeuron == null && nucleus.parent.instanceCount > 1) {
this.selectedSynapseNeuron = nucleus;
this.expandArray = false;
}
else {
this.currentNucleus = nucleus;
if (this.currentNucleus is Neuron neuron && neuron.receivers.Count == 0)
this.selectedOutput = this.currentNucleus;
this.selectedSynapseNeuron = null;
this.expandArray = false;
}
}
else {
// select the cluster, not the neuron in the cluster
this.currentNucleus = nucleus.parent;
this.expandArray = false;
}
}
else {
this.currentNucleus = nucleus;
if (this.currentNucleus is Neuron neuron && neuron.receivers.Count == 0)
this.selectedOutput = this.currentNucleus;
this.expandArray = false;
}
}
protected void OnClusterClick(Cluster subCluster) {
// May be used with storedPrefab...
Selection.activeObject = subCluster.prefab;
EditorGUIUtility.PingObject(subCluster.prefab);
ClusterViewer.previousPrefab = this.currentCluster.prefab;
Editor.CreateEditor(subCluster.prefab);
}
protected void OnAllOutputsClick() {
this.currentNucleus = null;
this.selectedOutput = null;
this.expandArray = false;
}
#endregion Interaction
}
}