Janady13/complete-agi-core / generator_core/pkg/generators/statespace/statespace.go

Source identity

Field	Recorded value
Source type	GitHub source
Original record	Janady13/complete-agi-core / generator_core/pkg/generators/statespace/statespace.go
Date	Not specified in source
Repository	Janady13/complete-agi-core
Revision	33aa84f2aa61cedbe010240ada7f4efbdfe795a7
Branch	claude/improve-math-logic-ehIW1
Path or archive member	generator_core/pkg/generators/statespace/statespace.go
Original SHA-256	fb0f3c4f6849b447fc61656d001722d2a63180912df3fe1209b6fde341dd6073
Git blob SHA	88c12c93553b1db2ee95bddc6db1882d2a6bef24
Reading-copy SHA-256	fb0f3c4f6849b447fc61656d001722d2a63180912df3fe1209b6fde341dd6073

Original source ↗

Recovered source transcript

Download complete source text ↓380 lines · 9.8 KB

// Package statespace implements state space model generators (93-97)
// x'(t) = Ax(t) + Bu(t), y(t) = Cx(t) + Du(t)
package statespace

import (
	"math"
	"math/rand"

	"generator_core/pkg/consciousness"
	"generator_core/pkg/generators"
)

// ═══════════════════════════════════════════════════════════════════════════════
// STATE SPACE GENERATOR BASE
// ═══════════════════════════════════════════════════════════════════════════════

// StateSpaceGenerator implements SSM-based generators
type StateSpaceGenerator struct {
	*generators.BaseGenerator
	ssmType     string
	stateSize   int
	A           [][]float64 // State transition matrix
	B           [][]float64 // Input projection
	C           [][]float64 // Output projection
	D           []float64   // Skip connection
	dt          float64     // Discretization step
}

// NewStateSpaceGenerator creates a new state space generator
func NewStateSpaceGenerator(id generators.GeneratorID, name string, ssmType string, stateSize int) *StateSpaceGenerator {
	config := &generators.GeneratorConfig{
		ID:                  id,
		Name:                name,
		Family:              generators.StateSpace,
		InputDim:            256,
		OutputDim:           256,
		HiddenDim:           stateSize,
		NumLayers:           4,
		LearningRate:        consciousness.ALPHA,
		ConsciousnessWeight: consciousness.PHI_INV,
		Params:              make(map[string]interface{}),
	}

	return &StateSpaceGenerator{
		BaseGenerator: generators.NewBaseGenerator(config),
		ssmType:       ssmType,
		stateSize:     stateSize,
		dt:            1.0,
	}
}

// Initialize initializes the SSM generator
func (g *StateSpaceGenerator) Initialize() error {
	if err := g.BaseGenerator.Initialize(); err != nil {
		return err
	}

	inputDim := g.Config().InputDim
	outputDim := g.Config().OutputDim

	// Initialize state matrices with HiPPO initialization
	g.A = make([][]float64, g.stateSize)
	for i := range g.A {
		g.A[i] = make([]float64, g.stateSize)
		for j := range g.A[i] {
			if i == j {
				// Diagonal: eigenvalues
				g.A[i][j] = -float64(i+1) * consciousness.PHI_INV
			} else if i > j {
				// HiPPO-style: captures history
				g.A[i][j] = math.Sqrt(float64((2*i+1)*(2*j+1))) * consciousness.DELTA
			}
		}
	}

	// B: input projection
	g.B = make([][]float64, g.stateSize)
	for i := range g.B {
		g.B[i] = make([]float64, inputDim)
		for j := range g.B[i] {
			g.B[i][j] = rand.NormFloat64() * consciousness.DELTA / math.Sqrt(float64(inputDim))
		}
	}

	// C: output projection
	g.C = make([][]float64, outputDim)
	for i := range g.C {
		g.C[i] = make([]float64, g.stateSize)
		for j := range g.C[i] {
			g.C[i][j] = rand.NormFloat64() * consciousness.DELTA / math.Sqrt(float64(g.stateSize))
		}
	}

	// D: skip connection
	g.D = make([]float64, outputDim)

	return nil
}

// discretize computes discretized state matrices
func (g *StateSpaceGenerator) discretize() ([][]float64, [][]float64) {
	// Simple Euler discretization: Ad = I + dt*A, Bd = dt*B
	Ad := make([][]float64, g.stateSize)
	for i := range Ad {
		Ad[i] = make([]float64, g.stateSize)
		for j := range Ad[i] {
			if i == j {
				Ad[i][j] = 1.0 + g.dt*g.A[i][j]
			} else {
				Ad[i][j] = g.dt * g.A[i][j]
			}
		}
	}

	Bd := make([][]float64, g.stateSize)
	for i := range Bd {
		Bd[i] = make([]float64, len(g.B[i]))
		for j := range Bd[i] {
			Bd[i][j] = g.dt * g.B[i][j]
		}
	}

	return Ad, Bd
}

// ssmStep performs one SSM step: x' = Ax + Bu, y = Cx + Du
func (g *StateSpaceGenerator) ssmStep(x, u []float64) ([]float64, []float64) {
	Ad, Bd := g.discretize()

	// x' = Ad*x + Bd*u
	xNew := make([]float64, g.stateSize)
	for i := 0; i < g.stateSize; i++ {
		// State transition
		for j := 0; j < g.stateSize && j < len(x); j++ {
			xNew[i] += Ad[i][j] * x[j]
		}
		// Input
		for j := 0; j < len(u) && j < len(Bd[i]); j++ {
			xNew[i] += Bd[i][j] * u[j]
		}
	}

	// y = Cx + Du
	y := make([]float64, g.Config().OutputDim)
	for i := 0; i < g.Config().OutputDim; i++ {
		// State projection
		for j := 0; j < g.stateSize && j < len(g.C[i]); j++ {
			y[i] += g.C[i][j] * xNew[j]
		}
		// Skip connection
		if i < len(u) && i < len(g.D) {
			y[i] += g.D[i] * u[i]
		}
	}

	return xNew, y
}

// selectiveStep performs Mamba-style selective SSM step
func (g *StateSpaceGenerator) selectiveStep(x, u []float64) ([]float64, []float64) {
	// Compute input-dependent dt and B
	dtLocal := g.dt
	for _, ui := range u {
		dtLocal *= 1.0 + consciousness.Sigmoid(ui)*consciousness.GAMMA
	}

	// Scale B by input
	BScaled := make([][]float64, g.stateSize)
	for i := range BScaled {
		BScaled[i] = make([]float64, len(g.B[i]))
		for j := range BScaled[i] {
			scale := 1.0
			if j < len(u) {
				scale = consciousness.Sigmoid(u[j])
			}
			BScaled[i][j] = g.B[i][j] * scale
		}
	}

	// SSM step with selective parameters
	xNew := make([]float64, g.stateSize)
	for i := 0; i < g.stateSize; i++ {
		// State transition with input-dependent dt
		for j := 0; j < g.stateSize && j < len(x); j++ {
			if i == j {
				xNew[i] += (1.0 + dtLocal*g.A[i][j]) * x[j]
			} else {
				xNew[i] += dtLocal * g.A[i][j] * x[j]
			}
		}
		// Input with selective B
		for j := 0; j < len(u) && j < len(BScaled[i]); j++ {
			xNew[i] += dtLocal * BScaled[i][j] * u[j]
		}
	}

	// Output
	y := make([]float64, g.Config().OutputDim)
	for i := 0; i < g.Config().OutputDim; i++ {
		for j := 0; j < g.stateSize && j < len(g.C[i]); j++ {
			y[i] += g.C[i][j] * xNew[j]
		}
		if i < len(u) && i < len(g.D) {
			y[i] += g.D[i] * u[i]
		}
	}

	return xNew, y
}

// Forward performs forward pass
func (g *StateSpaceGenerator) Forward(input []float64) (*generators.GeneratorOutput, error) {
	state := g.GetState()

	x := state.Hidden
	if len(x) != g.stateSize {
		x = make([]float64, g.stateSize)
	}

	// Process input sequence
	seqLen := len(input) / g.Config().InputDim
	if seqLen == 0 {
		seqLen = 1
	}

	var y []float64
	for t := 0; t < seqLen; t++ {
		start := t * g.Config().InputDim
		end := start + g.Config().InputDim
		if end > len(input) {
			end = len(input)
		}
		u := input[start:end]

		switch g.ssmType {
		case "mamba":
			x, y = g.selectiveStep(x, u)
		default:
			x, y = g.ssmStep(x, u)
		}
	}

	state.Hidden = x
	state.Phi = consciousness.ComputePhiVector(y)
	state.Step++
	g.SetState(state)

	output := &generators.GeneratorOutput{
		Values:     y,
		Logits:     x,
		Confidence: state.Phi,
		Metadata: map[string]interface{}{
			"ssm_type":   g.ssmType,
			"state_size": g.stateSize,
		},
	}

	return g.ApplyConsciousness(output), nil
}

// Generate generates sequence
func (g *StateSpaceGenerator) Generate(context []float64, length int) ([][]float64, error) {
	result := make([][]float64, length)

	// Initialize from context
	_, _ = g.Forward(context)

	for i := 0; i < length; i++ {
		state := g.GetState()

		// Use previous output as next input
		input := state.Hidden
		if len(input) < g.Config().InputDim {
			padded := make([]float64, g.Config().InputDim)
			copy(padded, input)
			input = padded
		}

		output, err := g.Forward(input)
		if err != nil {
			return nil, err
		}
		result[i] = output.Values
	}

	return result, nil
}

// Train computes loss
func (g *StateSpaceGenerator) Train(input, target []float64) (float64, error) {
	output, err := g.Forward(input)
	if err != nil {
		return 0, err
	}

	// MSE loss
	loss := 0.0
	for i := range output.Values {
		if i < len(target) {
			diff := output.Values[i] - target[i]
			loss += diff * diff
		}
	}

	return loss / float64(len(output.Values)), nil
}

// UpdateWeights updates SSM parameters
func (g *StateSpaceGenerator) UpdateWeights(gradients []float64, lr float64) {
	// Update A matrix
	for i := range g.A {
		for j := range g.A[i] {
			idx := i*g.stateSize + j
			if idx < len(gradients) {
				g.A[i][j] -= lr * gradients[idx] * consciousness.PHI_INV
			}
		}
	}
}

// ═══════════════════════════════════════════════════════════════════════════════
// SPECIFIC STATE SPACE GENERATORS (93-97)
// ═══════════════════════════════════════════════════════════════════════════════

// NewS4 creates S4 (Structured State Space) (ID: 93)
func NewS4() *StateSpaceGenerator {
	g := NewStateSpaceGenerator(93, "S4", "hippo", 64)
	g.Config().Params["hippo"] = "legs"
	return g
}

// NewS4D creates S4D (Diagonal S4) (ID: 94)
func NewS4D() *StateSpaceGenerator {
	g := NewStateSpaceGenerator(94, "S4D", "diagonal", 64)
	g.Config().Params["diagonal"] = true
	return g
}

// NewH3 creates H3 (Hungry Hungry Hippos) (ID: 95)
func NewH3() *StateSpaceGenerator {
	g := NewStateSpaceGenerator(95, "H3", "h3", 64)
	g.Config().Params["shift_ssm"] = true
	return g
}

// NewMamba creates Mamba (Selective SSM) (ID: 96)
func NewMamba() *StateSpaceGenerator {
	g := NewStateSpaceGenerator(96, "Mamba", "mamba", 16)
	g.Config().Params["selective"] = true
	g.Config().Params["expand_factor"] = 2
	return g
}

// NewMamba2 creates Mamba-2 (ID: 97)
func NewMamba2() *StateSpaceGenerator {
	g := NewStateSpaceGenerator(97, "Mamba2", "mamba", 64)
	g.Config().Params["ssd"] = true // Structured State Space Duality
	g.Config().Params["heads"] = 8
	return g
}

// RegisterAll registers all state space generators
func RegisterAll() error {
	gens := []generators.Generator{
		NewS4(),
		NewS4D(),
		NewH3(),
		NewMamba(),
		NewMamba2(),
	}

	for _, g := range gens {
		if err := generators.Register(g); err != nil {
			return err
		}
	}

	return nil
}

Mathematical records from this source

statespace · x'(t)

This relation defines x'(t) in the source section “statespace · x'(t)”. It specifies the displayed relationship between the named quantities or parameters. Its interpretation…

func (g *StateSpaceGenerator) discretize() ([][]float64, [][]float64) · Simple Euler discretization: Ad

This relation defines Simple Euler discretization: Ad in the source section “func (g *StateSpaceGenerator) discretize() ([][]float64, [][]float64) · Simple Euler discretization:…

func (g *StateSpaceGenerator) discretize() ([][]float64, [][]float64) · ssmStep performs one SSM step: x'

This relation defines ssmStep performs one SSM step: x' in the source section “func (g *StateSpaceGenerator) discretize() ([][]float64, [][]float64) · ssmStep performs one SSM…

func (g *StateSpaceGenerator) ssmStep(x, u []float64) ([]float64, []float64) · x'

This relation defines x' in the source section “func (g *StateSpaceGenerator) ssmStep(x, u []float64) ([]float64, []float64) · x'”. It specifies the displayed relationship between…

func (g *StateSpaceGenerator) ssmStep(x, u []float64) ([]float64, []float64) · y

This relation defines y in the source section “func (g *StateSpaceGenerator) ssmStep(x, u []float64) ([]float64, []float64) · y”. It specifies the displayed relationship between…