diff --git a/@tree/treefun.m b/@tree/treefun.m index f93c186..c2a8d43 100644 --- a/@tree/treefun.m +++ b/@tree/treefun.m @@ -1,16 +1,36 @@ -function newTree = treefun(obj, fun) +function varargout = treefun(obj, fun) %% TREEFUN Create a new tree by applying function handle to each node of the source tree. - if ~isa(fun, 'function_handle') error('MATLAB:tree:treefun', ... 'Second argument, fun, must be a function handle. Got a %s.', ... class(fun)); end - - % First we copy - newTree = tree(obj, 'clear'); - - % Then we override - newTree.Node = cellfun(fun, obj.Node, 'UniformOutput', false); - + % validate number of output arguments is <= # outputs of `fun` + nargoutchk(0,abs(nargout(fun))); + + % if called interactively, nargout=0, which causes problems in `cellfun` + % given: + % >> t + % t = ... + % Node: {3x1 cell} + % e.g. if this isn't done... + % >> t.treefun(@(x) x+1) + % >> + % but by setting output number to 1, `ans` will be assigned + % >> t.treefun(@(x) x+1); + % ans = ... + % Node: {3x1 cell} + if ~nargout + n_out = 1; % need new variable b/c can't + % reassign nargout + else + n_out = nargout; + end + + % cellfun call collects output into n_outx1 cell array + [nodes{1:n_out}] = cellfun(fun,obj.Node,'UniformOutput',false); + for k = 1:n_out % copy each output to its own tree + varargout{k} = tree(obj,'clear'); + varargout{k}.Node = nodes{k}; + end end \ No newline at end of file diff --git a/@tree/treefun2.m b/@tree/treefun2.m index ed36079..b65d605 100644 --- a/@tree/treefun2.m +++ b/@tree/treefun2.m @@ -1,19 +1,32 @@ -function obj = treefun2(obj, val, fun) +function varargout = treefun2(obj, val, fun) %%TREEFUN2 Two-arguments function on trees, with scalar expansion. + if ~isa(fun, 'function_handle') + error('MATLAB:tree:treefun', ... + 'Second argument, fun, must be a function handle. Got a %s.', ... + class(fun)); + end + nargoutchk(0,abs(nargout(fun))); % validate # output args [obj, val] = permuteIfNeeded(obj, val); - if isa(val, 'tree') - - content = cellfun(fun, obj.Node, val.Node, ... - 'UniformOutput', false); + % if nargout = 0, set it to 1 so that `ans` is assigned + % (see comment in `treefun.m`) + if ~nargout, n_out = 1; else n_out = nargout; end + + if isa(val, 'tree') - else - - content = cellfun(@(x) fun(x, val), obj.Node, ... - 'UniformOutput', false); + [content{1:n_out}] = cellfun(fun, obj.Node, val.Node,... + 'UniformOutput',false); - end + else + + [content{1:n_out}] = cellfun(@(x) fun(x, val), obj.Node, ... + 'UniformOutput', false); + + end - obj.Node = content; + for k = 1:n_out + varargout{k} = tree(obj,'clear'); + varargout{k}.Node = content{k}; + end end \ No newline at end of file