Expandir / recolher a árvore de materiais usando dados booleanos (Angular 7)

Estou tentando obter uma árvore de materiais trabalhando com booleanos de sua fonte de dados. Estou filtrando datasource.data na entrada, o que causa o colapso da árvore sempre que o usuário insere algo. Para evitar isso, gostaria de usar um booleano na linha específica para expandir / recolher a árvore de materiais. Fiz isso com sucesso com a mesa de material, mas não consigo fazer funcionar com a árvore. Isso é o que eu tenho até agora:

Arquivo HTML:

<mat-tree [dataSource]="nestedDataSource" [treeControl]="nestedTreeControl" class="example-tree">
  <mat-tree-node *matTreeNodeDef="let node">
    <li class="mat-tree-node">
      <button mat-icon-button disabled></button>
      {{node.name}}
    </li>
  </mat-tree-node>

  <mat-nested-tree-node *matTreeNodeDef="let node; when: hasNestedChild">
    <li>
      <div class="mat-tree-node">
        <button mat-icon-button
                [attr.aria-label]="'toggle ' + node.name"
                click="changeState(node)">
          <mat-icon class="mat-icon-rtl-mirror">
            {{node.expanded ? 'expand_more' : 'chevron_right'}}
          </mat-icon>
        </button>
        {{node.name}} ({{node.expanded}})
      </div>
      <ul [class.example-tree-invisible]="node.expanded">
        <ng-container matTreeNodeOutlet></ng-container>
      </ul>
    </li>
  </mat-nested-tree-node>
</mat-tree>

Arquivo TS:

import {NestedTreeControl} from '@angular/cdk/tree';
import {Component, Injectable} from '@angular/core';
import {MatTreeNestedDataSource} from '@angular/material/tree';
import {BehaviorSubject} from 'rxjs';

/**
 * Json node data with nested structure. Each node has a filename and a value or a list of children
 */
export class FileNode {
  childGroups: FileNode[];
  name: string;
  expanded: boolean;
}

/**
 * The Json tree data in string. The data could be parsed into Json object
 */
const TREE_DATA = [
  {
  'name': 'Group 1',
  'expanded': false,
  'childGroups': [
    {
      'name': 'Childgroup 1',
      'expanded': false,
      'childGroups': []
    },
    {
      'name': 'Childgroup 2',
      'expanded': false,
      'childGroups': [
        {
          'name': 'Child of child',
          'expanded': false,
          'childGroups': []
        }
      ]
    }
    ]
  },
    {
  'name': 'Group 2',
  'expanded': false,
  'childGroups': [
    {
      'name': 'Childgroup 1',
      'expanded': false,
      'childGroups': []
    },
    {
      'name': 'Childgroup 2',
      'expanded': false,
      'childGroups': [
        {
          'name': 'Child of child',
          'expanded': false,
          'childGroups': []
        }
      ]
    }
    ]
  },
    {
  'name': 'Group 3',
  'expanded': false,
  'childGroups': [
    {
      'name': 'Childgroup 1',
      'expanded': false,
      'childGroups': []
    },
    {
      'name': 'Childgroup 2',
      'expanded': false,
      'childGroups': [
        {
          'name': 'Child of child',
          'expanded': false,
          'childGroups': []
        }
      ]
    }
    ]
  },
]

@Component({
  selector: 'tree-nested-overview-example',
  templateUrl: 'tree-nested-overview-example.html',
  styleUrls: ['tree-nested-overview-example.css']
})
export class TreeNestedOverviewExample {
  nestedTreeControl: NestedTreeControl<FileNode>;
  nestedDataSource: MatTreeNestedDataSource<FileNode>;

  constructor() {
    this.nestedTreeControl = new NestedTreeControl<FileNode>(this._getChildren);
    this.nestedDataSource = new MatTreeNestedDataSource();
    this.nestedDataSource.data = TREE_DATA;
  }

  hasNestedChild = (_: number, nodeData: FileNode) => 
  nodeData.childGroups.length > 0;

  private _getChildren = (node: FileNode) => node.childGroups;

  changeState(node) {
    console.log(node);
    node.expanded = !node.expanded;
  }
}

Stackblitz